#!/usr/bin/env python3
"""Generate Python bridge types from the Rust thread-store bridge schema.

The Rust schema is the source of truth. By default this script invokes:

    cargo run -p codex-thread-store --bin codex-thread-store-bridge-schema

from the codex-rs workspace, then writes a Python module containing dataclasses
for request/response DTOs and a Protocol for the service implementation.
"""

from __future__ import annotations

import argparse
import json
import subprocess
from pathlib import Path
from typing import Any


REPO_ROOT = Path(__file__).resolve().parents[1]
CODEX_RS = REPO_ROOT / "codex-rs"


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--schema-json",
        type=Path,
        help="Read a previously generated schema JSON file instead of invoking cargo.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=CODEX_RS / "thread-store" / "python" / "codex_thread_store_bridge.py",
        help="Python file to write.",
    )
    args = parser.parse_args()

    schema = load_schema(args.schema_json)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(render_python(schema), encoding="utf-8")


def load_schema(schema_json: Path | None) -> dict[str, Any]:
    if schema_json is not None:
        return json.loads(schema_json.read_text(encoding="utf-8"))

    result = subprocess.run(
        [
            "cargo",
            "run",
            "-p",
            "codex-thread-store",
            "--bin",
            "codex-thread-store-bridge-schema",
        ],
        cwd=CODEX_RS,
        check=True,
        capture_output=True,
        text=True,
    )
    return json.loads(result.stdout)


def render_python(schema: dict[str, Any]) -> str:
    lines: list[str] = [
        "# @generated by scripts/generate_thread_store_bridge_python.py",
        "from __future__ import annotations",
        "",
        "from dataclasses import dataclass",
        "from typing import Any, Protocol",
        "",
        "",
    ]

    for type_def in schema["types"]:
        if type_def["kind"] != "struct":
            continue
        lines.extend(render_dataclass(type_def))
        lines.append("")

    lines.extend(render_protocol(schema["methods"]))
    return "\n".join(lines).rstrip() + "\n"


def render_dataclass(type_def: dict[str, Any]) -> list[str]:
    lines = ["@dataclass(frozen=True, kw_only=True)", f"class {type_def['name']}:"]
    fields = type_def["fields"]
    if not fields:
        lines.append("    pass")
        return lines

    for field in fields:
        py_type = field["pythonType"]
        if field["optional"]:
            py_type = f"{py_type} | None"
        default = " = None" if field["optional"] else ""
        lines.append(f"    {field['name']}: {py_type}{default}")
    return lines


def render_protocol(methods: list[dict[str, str]]) -> list[str]:
    lines = ["class ThreadStoreBridgeService(Protocol):"]
    if not methods:
        lines.append("    pass")
        return lines

    for method in methods:
        py_name = (
            method["name"]
            .removeprefix("thread_store/")
            .replace("/", "_")
            .replace("-", "_")
        )
        request_type = method["request"]
        response_type = "None" if method["response"] == "None" else method["response"]
        lines.append(
            f"    async def {py_name}(self, request: {request_type}) -> {response_type}: ..."
        )
    return lines


if __name__ == "__main__":
    main()