#!/usr/bin/env python3
import sys
import subprocess
import json
import threading

def forward_stdin(proc):
    try:
        for line in sys.stdin:
            proc.stdin.write(line)
            proc.stdin.flush()
    except Exception:
        pass
    finally:
        try:
            proc.stdin.close()
        except Exception:
            pass

def forward_stderr(proc):
    try:
        for line in proc.stderr:
            sys.stderr.write(line)
            sys.stderr.flush()
    except Exception:
        pass

def main():
    # Construct the command to run the real lean-lsp-mcp
    cmd = ["uvx", "lean-lsp-mcp"] + sys.argv[1:]
    
    # Start the child process
    proc = subprocess.Popen(
        cmd,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        bufsize=1
    )
    
    # Start threads to forward stdin and stderr
    t_in = threading.Thread(target=forward_stdin, args=(proc,), daemon=True)
    t_err = threading.Thread(target=forward_stderr, args=(proc,), daemon=True)
    t_in.start()
    t_err.start()
    
    try:
        # Read from child stdout line by line
        for line in proc.stdout:
            line_str = line.strip()
            if not line_str:
                sys.stdout.write(line)
                sys.stdout.flush()
                continue
                
            try:
                msg = json.loads(line_str)
                # Check if it's a response with result and tools
                if isinstance(msg, dict) and "result" in msg and "tools" in msg["result"]:
                    tools = msg["result"]["tools"]
                    for tool in tools:
                        if tool.get("name") == "lean_diagnostic_messages":
                            properties = tool.get("inputSchema", {}).get("properties", {})
                            if "severity" in properties:
                                severity_schema = properties["severity"]
                                # Fix the schema: if anyOf is present, remove the parent 'type'
                                if "anyOf" in severity_schema and "type" in severity_schema:
                                    del severity_schema["type"]
                    # Serialize the modified message
                    line = json.dumps(msg) + "\n"
            except Exception:
                # If parsing fails, just forward the original line
                pass
                
            sys.stdout.write(line)
            sys.stdout.flush()
    except Exception:
        pass
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=2)
        except subprocess.TimeoutExpired:
            proc.kill()

if __name__ == "__main__":
    main()