"""Project-bound, task-owned MCP session. Execute reviewed native UE commands."""
import argparse
import asyncio
import base64
import json
import logging
import sys
from pathlib import Path
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client
from mcp.types import Implementation
from CheckMcpReadiness import ROOT, ENDPOINT, verify_owner, unpack, redact, Warnings


async def run(args):
    owner = verify_owner()
    warnings = []
    logger = logging.getLogger("mcp.client.streamable_http")
    handler = Warnings(warnings)
    logger.addHandler(handler)
    logger.propagate = False
    commands = json.loads(Path(args.file).read_text(encoding="utf-8-sig"))
    if isinstance(commands, dict):
        commands = [commands]
    output = {"project": owner, "commands": [], "warnings": warnings}
    try:
        async with streamablehttp_client(ENDPOINT, timeout=180, sse_read_timeout=240, terminate_on_close=True) as (r, w, _):
            async with ClientSession(r, w, client_info=Implementation(name="irons-migration", version="1.0")) as session:
                await session.initialize()
                async def call(toolset, name, arguments=None):
                    return unpack(await session.call_tool("call_tool", {"toolset_name":toolset, "tool_name":name,"arguments":arguments or {}}))
                output["pieBefore"] = await call("EditorToolset.EditorAppToolset", "IsPIERunning")
                output["levelBefore"] = await call("IronsEditor.IronsMigrationToolset", "Snapshot") if output["pieBefore"].get("returnValue") else await call("editor_toolset.toolsets.scene.SceneTools", "get_current_level")
                for command in commands:
                    if verify_owner()["processId"] != owner["processId"]:
                        raise RuntimeError("Target process changed")
                    method = command.get("method", "call_tool")
                    arguments = command.get("arguments", {})
                    payload = arguments if method != "call_tool" else {k: command[k] for k in ("toolset_name","tool_name","arguments") if k in command}
                    result = await session.call_tool(method, payload)
                    data = unpack(result)
                    if "image_path" in command:
                        payload = data.get("returnValue", {})
                        image = payload.get("image",payload)
                        encoded = image.get("data") or next((b.data for b in result.content if b.type == "image"), None)
                        if not encoded:
                            raise RuntimeError("Capture did not return pixels")
                        path = Path(command["image_path"]).resolve()
                        path.relative_to(ROOT)
                        path.parent.mkdir(parents=True, exist_ok=True)
                        path.write_bytes(base64.b64decode(encoded))
                        data = {"savedImage": str(path)}
                    output["commands"].append({"tool": command.get("tool_name", method),"result":data})
                    if args.print_results:
                        print(redact(json.dumps(output["commands"][-1], ensure_ascii=False)), flush=True)
                output["pieAfter"] = await call("EditorToolset.EditorAppToolset", "IsPIERunning")
                output["levelAfter"] = await call("IronsEditor.IronsMigrationToolset", "Snapshot") if output["pieAfter"].get("returnValue") else await call("editor_toolset.toolsets.scene.SceneTools", "get_current_level")
        if verify_owner()["processId"] != owner["processId"]:
            raise RuntimeError("Target process changed")
        output["clientClosed"] = True
    finally:
        logger.removeHandler(handler)
        path = ROOT / args.output
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(redact(json.dumps(output, ensure_ascii=False, indent=2))+"\n",encoding="utf-8")
    print(json.dumps({"record":str(path),"calls":len(output["commands"]),"warnings":warnings},ensure_ascii=False))


if __name__ == "__main__":
    sys.stdout.reconfigure(encoding="utf-8")
    parser = argparse.ArgumentParser()
    parser.add_argument("--file",required=True)
    parser.add_argument("--output",default="Saved/Migration/mcp-results.json")
    parser.add_argument("--print-results",action="store_true")
    asyncio.run(run(parser.parse_args()))