#!/usr/bin/env python3
"""Code generator for zuzu SDK components.

Reads JSON schema files from schemas/ and generates:
  - zuzu/components/_generated/{category}.py   (Python functions)
  - zuzu/components/_generated/__init__.py      (package init)
  - zuzu/_generated_exports.py                  (all imports + __all__)
  - frontend/src/types/_generated/components.ts (TypeScript interfaces)
  - tests/test_components_generated.py          (spec snapshot tests)

Usage:
  python3 scripts/codegen.py              # Generate all files
  python3 scripts/codegen.py --dry-run    # Print summary without writing
  python3 scripts/codegen.py --check      # Exit 1 if generated files are stale
"""
from __future__ import annotations

import hashlib
import json
import sys
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parent.parent
SCHEMAS_DIR = ROOT / "schemas"
GENERATED_PY_DIR = ROOT / "zuzu" / "components" / "_generated"
GENERATED_TS_DIR = ROOT / "frontend" / "src" / "types" / "_generated"
EXPORTS_FILE = ROOT / "zuzu" / "_generated_exports.py"
TESTS_FILE = ROOT / "tests" / "test_components_generated.py"

# ── Type mappings ────────────────────────────────────────────────────────

PY_TYPE_MAP: dict[str, str] = {
    "string": "str",
    "number": "int | float",
    "boolean": "bool",
    "array": "list[Any]",
    "object": "dict[str, Any]",
    "any": "Any",
}

TS_TYPE_MAP: dict[str, str] = {
    "string": "string",
    "number": "number",
    "boolean": "boolean",
    "array": "any[]",
    "object": "Record<string, any>",
    "any": "any",
}


# ── Schema loading ───────────────────────────────────────────────────────

def load_meta() -> dict[str, Any]:
    meta_path = SCHEMAS_DIR / "_meta.json"
    if not meta_path.exists():
        print(f"ERROR: {meta_path} not found", file=sys.stderr)
        sys.exit(1)
    with open(meta_path) as f:
        return json.load(f)


def load_schemas(categories: list[str]) -> dict[str, list[dict[str, Any]]]:
    """Load all schemas grouped by category."""
    grouped: dict[str, list[dict[str, Any]]] = {c: [] for c in categories}
    for cat in categories:
        cat_dir = SCHEMAS_DIR / cat
        if not cat_dir.is_dir():
            continue
        for path in sorted(cat_dir.glob("*.json")):
            with open(path) as f:
                schema = json.load(f)
            _validate_schema(schema, path)
            grouped[cat].append(schema)
    return grouped


def _validate_schema(schema: dict[str, Any], path: Path) -> None:
    required = {"name", "category", "mode", "spec_type", "doc", "props"}
    missing = required - set(schema.keys())
    if missing:
        print(f"ERROR: {path} missing fields: {missing}", file=sys.stderr)
        sys.exit(1)
    for prop in schema["props"]:
        if "name" not in prop or "type" not in prop:
            print(f"ERROR: {path} prop missing name/type: {prop}", file=sys.stderr)
            sys.exit(1)


# ── Python generation ────────────────────────────────────────────────────

def _py_header(category: str) -> str:
    lines = [
        '"""AUTO-GENERATED by scripts/codegen.py \u2014 DO NOT EDIT."""',
        "from __future__ import annotations",
        "",
        "from typing import Any, Callable",
        "",
        "from zuzu.runtime.spec import _add_component",
        "from zuzu.record import Record, RecordQuery",
        "from zuzu.components._utils import _check_visible",
    ]
    return "\n".join(lines) + "\n"


def _py_type(prop: dict[str, Any]) -> str:
    """Return Python type hint for a prop."""
    if "python_type" in prop:
        return prop["python_type"]
    base = PY_TYPE_MAP.get(prop["type"], "Any")
    if not prop.get("required", False) and prop.get("default") is None:
        return f"{base} | None"
    return base


def _py_default(prop: dict[str, Any]) -> str:
    """Return Python default value literal."""
    d = prop.get("default")
    if d is None:
        return "None"
    if isinstance(d, bool):
        return "True" if d else "False"
    if isinstance(d, str):
        return repr(d)
    if isinstance(d, (int, float)):
        return str(d)
    if isinstance(d, list):
        return repr(d)
    if isinstance(d, dict):
        return repr(d)
    return repr(d)


def _prop_value_expr(prop: dict[str, Any]) -> str:
    """Return the expression to use for the prop value in spec dict."""
    if "python_transform" in prop:
        return prop["python_transform"]
    return prop["name"]


def _spec_key(prop: dict[str, Any]) -> str:
    """Return the key name for the spec dict."""
    return prop.get("spec_key", prop["name"])


def _is_always_included(prop: dict[str, Any]) -> bool:
    """Props with include_if='always' or that are required+positional."""
    return prop.get("include_if") == "always"


def gen_python_function(schema: dict[str, Any]) -> str:
    """Generate a single Python component function."""
    mode = schema["mode"]
    name = schema["name"]
    spec_type = schema["spec_type"]
    doc = schema["doc"]
    props = schema["props"]

    # Split props into positional and keyword-only
    positional: list[dict[str, Any]] = []
    keyword: list[dict[str, Any]] = []
    for p in props:
        if p.get("positional", False):
            positional.append(p)
        else:
            keyword.append(p)

    # Build parameter list
    params: list[str] = []

    # For data_bound, first param is source
    if mode == "data_bound":
        params.append("source: Record | RecordQuery | list[dict[str, Any]]")

    # Positional params (required first, then optional)
    pos_required = [p for p in positional if p.get("required", False)]
    pos_optional = [p for p in positional if not p.get("required", False)]

    for p in pos_required:
        params.append(f"{p['name']}: {_py_type(p)}")
    for p in pos_optional:
        params.append(f"{p['name']}: {_py_type(p)} = {_py_default(p)}")

    # Keyword-only separator
    kw_params: list[str] = []
    for p in keyword:
        if p.get("required", False):
            kw_params.append(f"{p['name']}: {_py_type(p)}")
        else:
            kw_params.append(f"{p['name']}: {_py_type(p)} = {_py_default(p)}")

    # Auto-inject visible for append and data_bound modes
    if mode in ("append", "data_bound"):
        kw_params.append("visible: Any = True")

    # Build function signature
    sig_parts: list[str] = []
    sig_parts.extend(params)
    if kw_params:
        sig_parts.append("*")
        sig_parts.extend(kw_params)

    # Format signature
    sig = ", ".join(sig_parts)

    # Return type
    if mode == "return":
        ret_type = "dict[str, Any]"
    else:
        ret_type = "None"

    # Determine which props are always included in the base dict
    always_props: list[dict[str, Any]] = []
    conditional_props: list[dict[str, Any]] = []
    for p in props:
        if p.get("required", False) or _is_always_included(p):
            always_props.append(p)
        else:
            conditional_props.append(p)

    # Build function body
    lines: list[str] = []
    lines.append("")
    lines.append("")

    # Function def with wrapped params for readability
    if len(sig) > 70 or len(sig_parts) > 3:
        lines.append(f"def {name}(")
        for i, part in enumerate(sig_parts):
            comma = "," if i < len(sig_parts) - 1 else ","
            lines.append(f"    {part}{comma}")
        lines.append(f") -> {ret_type}:")
    else:
        lines.append(f"def {name}({sig}) -> {ret_type}:")

    # Docstring
    lines.append(f'    """{doc}"""')

    # Visible check for append/data_bound
    if mode in ("append", "data_bound"):
        lines.append("    if not _check_visible(visible):")
        lines.append("        return")

    # Variable name: 'comp' for append/data_bound, 'field' for return
    var = "comp" if mode in ("append", "data_bound") else "field"

    # Base dict construction
    base_items: list[str] = [f'"type": "{spec_type}"']
    for p in always_props:
        key = _spec_key(p)
        val = _prop_value_expr(p)
        base_items.append(f'"{key}": {val}')

    base_dict = ", ".join(base_items)
    lines.append(f"    {var}: dict[str, Any] = {{{base_dict}}}")

    # Source resolution for data_bound
    if mode == "data_bound":
        source_res = schema.get("source_resolution", "full")
        lines.append("")
        lines.append("    if isinstance(source, Record):")
        lines.append('        comp["source"] = source._table_name')
        lines.append("    elif isinstance(source, RecordQuery):")
        if source_res == "full":
            lines.append("        sql, params = source._query.build_select()")
            lines.append('        comp["source"] = source._record._table_name')
            lines.append('        comp["query"] = {"sql": sql, "params": params}')
        else:
            lines.append('        comp["source"] = source._record._table_name')
        lines.append("    else:")
        lines.append('        comp["data"] = source')

    # Conditional prop inclusion
    for p in conditional_props:
        key = _spec_key(p)
        val = _prop_value_expr(p)
        include_if = p.get("include_if", "truthy")

        if include_if == "not_none":
            lines.append(f"    if {p['name']} is not None:")
        else:
            # Default: truthy
            lines.append(f"    if {p['name']}:")
        lines.append(f'        {var}["{key}"] = {val}')

    # Final action
    if mode in ("append", "data_bound"):
        lines.append(f"    _add_component({var})")
    else:
        lines.append(f"    return {var}")

    return "\n".join(lines)


def gen_python_category(category: str, schemas: list[dict[str, Any]]) -> str:
    """Generate a full Python module for a category."""
    parts = [_py_header(category)]
    for schema in schemas:
        if schema["mode"] == "manual":
            continue
        parts.append(gen_python_function(schema))
    return "\n".join(parts) + "\n"


def gen_python_init(grouped: dict[str, list[dict[str, Any]]]) -> str:
    """Generate __init__.py for the _generated package."""
    lines = ['"""AUTO-GENERATED by scripts/codegen.py \u2014 DO NOT EDIT."""']
    for cat, schemas in grouped.items():
        names = [s["name"] for s in schemas if s["mode"] != "manual"]
        if names:
            names_str = ", ".join(names)
            lines.append(f"from zuzu.components._generated.{cat} import {names_str}")
    lines.append("")
    return "\n".join(lines) + "\n"


def gen_exports(grouped: dict[str, list[dict[str, Any]]]) -> str:
    """Generate zuzu/_generated_exports.py with all imports + __all__."""
    lines = ['"""AUTO-GENERATED by scripts/codegen.py \u2014 DO NOT EDIT."""']
    all_names: list[str] = []
    for cat, schemas in grouped.items():
        names = [s["name"] for s in schemas if s["mode"] != "manual"]
        if names:
            names_str = ", ".join(names)
            lines.append(f"from zuzu.components._generated.{cat} import {names_str}")
            all_names.extend(names)
    lines.append("")
    all_str = ", ".join(f'"{n}"' for n in sorted(all_names))
    lines.append(f"__all__ = [{all_str}]")
    lines.append("")
    return "\n".join(lines) + "\n"


# ── TypeScript generation ────────────────────────────────────────────────

def _ts_type(prop: dict[str, Any]) -> str:
    """Return TypeScript type for a prop."""
    base = TS_TYPE_MAP.get(prop["type"], "any")
    if prop.get("enum"):
        base = " | ".join(f'"{v}"' for v in prop["enum"])
    if not prop.get("required", False) and prop.get("default") is None:
        return f"{base} | null"
    return base


def _pascal_case(name: str) -> str:
    """Convert snake_case to PascalCase."""
    return "".join(word.capitalize() for word in name.split("_"))


def gen_typescript(all_schemas: list[dict[str, Any]]) -> str:
    """Generate TypeScript interfaces + union type."""
    lines = ["/* AUTO-GENERATED by scripts/codegen.py \u2014 DO NOT EDIT. */", ""]

    interface_names: list[str] = []

    for schema in all_schemas:
        if schema["mode"] == "manual":
            continue
        pascal = _pascal_case(schema["name"]) + "Spec"
        interface_names.append(pascal)

        lines.append(f"export interface {pascal} {{")
        lines.append(f'  type: "{schema["spec_type"]}";')

        for prop in schema["props"]:
            ts = _ts_type(prop)
            optional = "" if prop.get("required", False) else "?"
            key = prop.get("spec_key", prop["name"])
            lines.append(f"  {key}{optional}: {ts};")

        # For append/data_bound, add visible
        if schema["mode"] in ("append", "data_bound"):
            lines.append("  visible?: any;")

        # For data_bound, add source/data/query
        if schema["mode"] == "data_bound":
            lines.append("  source?: string;")
            lines.append("  data?: Record<string, any>[];")
            if schema.get("source_resolution") == "full":
                lines.append("  query?: { sql: string; params: any[] };")

        lines.append("}")
        lines.append("")

    # Union type
    if interface_names:
        union = " | ".join(interface_names)
        lines.append(f"export type AnyComponentSpec = {union};")
        lines.append("")

    return "\n".join(lines) + "\n"


# ── Test generation ──────────────────────────────────────────────────────

def gen_tests(all_schemas: list[dict[str, Any]]) -> str:
    """Generate spec snapshot tests."""
    lines = [
        '"""AUTO-GENERATED by scripts/codegen.py \u2014 DO NOT EDIT."""',
        "from __future__ import annotations",
        "",
        "import pytest",
        "from zuzu.runtime.spec import _components, _reset",
        "",
        "",
        '@pytest.fixture(autouse=True)',
        "def _clean():",
        "    _reset()",
        "    yield",
        "    _reset()",
        "",
    ]

    for schema in all_schemas:
        if schema["mode"] == "manual":
            continue

        name = schema["name"]
        spec_type = schema["spec_type"]
        cat = schema["category"]

        lines.append("")
        lines.append(f"def test_{name}_spec_type():")
        lines.append(f"    from zuzu.components._generated.{cat} import {name}")

        # Build minimal call with only required args
        args: list[str] = []
        kwargs: list[str] = []

        if schema["mode"] == "data_bound":
            args.append('[{"id": 1}]')

        for prop in schema["props"]:
            if not prop.get("required", False):
                continue
            val = _test_value(prop)
            if prop.get("positional", False):
                args.append(val)
            else:
                kwargs.append(f"{prop['name']}={val}")

        call_args = ", ".join(args + kwargs)

        if schema["mode"] == "return":
            lines.append(f"    result = {name}({call_args})")
            lines.append(f'    assert result["type"] == "{spec_type}"')
        else:
            lines.append(f"    {name}({call_args})")
            lines.append(f"    assert len(_components) == 1")
            lines.append(f'    assert _components[0]["type"] == "{spec_type}"')

    lines.append("")
    return "\n".join(lines) + "\n"


def _test_value(prop: dict[str, Any]) -> str:
    """Return a test value literal for a required prop."""
    t = prop["type"]
    if prop.get("enum"):
        return repr(prop["enum"][0])
    if t == "string":
        return repr("test")
    if t == "number":
        return "1"
    if t == "boolean":
        return "True"
    if t == "array":
        return "[]"
    if t == "object":
        return "{}"
    return '""'


# ── File I/O ─────────────────────────────────────────────────────────────

def _file_hash(path: Path) -> str | None:
    """Return SHA-256 of a file, or None if it doesn't exist."""
    if not path.exists():
        return None
    return hashlib.sha256(path.read_bytes()).hexdigest()


def _write_if_changed(path: Path, content: str, dry_run: bool) -> bool:
    """Write content to path if it differs. Returns True if changed."""
    new_hash = hashlib.sha256(content.encode()).hexdigest()
    old_hash = _file_hash(path)
    if new_hash == old_hash:
        return False
    if not dry_run:
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(content)
    return True


# ── Main ─────────────────────────────────────────────────────────────────

def main() -> None:
    dry_run = "--dry-run" in sys.argv
    check = "--check" in sys.argv

    meta = load_meta()
    categories = meta["categories"]
    grouped = load_schemas(categories)

    # Flatten for TS / tests
    all_schemas: list[dict[str, Any]] = []
    for cat in categories:
        all_schemas.extend(grouped[cat])

    total_schemas = len(all_schemas)
    non_manual = [s for s in all_schemas if s["mode"] != "manual"]
    print(f"Loaded {total_schemas} schemas ({len(non_manual)} code-generated, "
          f"{total_schemas - len(non_manual)} manual)")

    files_changed: list[str] = []
    files_unchanged: list[str] = []

    # 1. Python per category
    for cat in categories:
        schemas = grouped[cat]
        codegen_schemas = [s for s in schemas if s["mode"] != "manual"]
        if not codegen_schemas:
            continue
        content = gen_python_category(cat, schemas)
        path = GENERATED_PY_DIR / f"{cat}.py"
        if _write_if_changed(path, content, dry_run or check):
            files_changed.append(str(path.relative_to(ROOT)))
        else:
            files_unchanged.append(str(path.relative_to(ROOT)))

    # 2. __init__.py
    content = gen_python_init(grouped)
    path = GENERATED_PY_DIR / "__init__.py"
    if _write_if_changed(path, content, dry_run or check):
        files_changed.append(str(path.relative_to(ROOT)))
    else:
        files_unchanged.append(str(path.relative_to(ROOT)))

    # 3. _generated_exports.py
    content = gen_exports(grouped)
    if _write_if_changed(EXPORTS_FILE, content, dry_run or check):
        files_changed.append(str(EXPORTS_FILE.relative_to(ROOT)))
    else:
        files_unchanged.append(str(EXPORTS_FILE.relative_to(ROOT)))

    # 4. TypeScript
    content = gen_typescript(all_schemas)
    path = GENERATED_TS_DIR / "components.ts"
    if _write_if_changed(path, content, dry_run or check):
        files_changed.append(str(path.relative_to(ROOT)))
    else:
        files_unchanged.append(str(path.relative_to(ROOT)))

    # 5. Tests
    content = gen_tests(all_schemas)
    if _write_if_changed(TESTS_FILE, content, dry_run or check):
        files_changed.append(str(TESTS_FILE.relative_to(ROOT)))
    else:
        files_unchanged.append(str(TESTS_FILE.relative_to(ROOT)))

    # Summary
    print()
    if files_changed:
        verb = "Would write" if (dry_run or check) else "Wrote"
        for f in files_changed:
            print(f"  {verb}: {f}")
    if files_unchanged:
        for f in files_unchanged:
            print(f"  Unchanged: {f}")

    print(f"\n{len(files_changed)} changed, {len(files_unchanged)} unchanged")

    if check and files_changed:
        print("\nERROR: Generated files are stale. Run: python3 scripts/codegen.py")
        sys.exit(1)

    if dry_run:
        print("\n(dry-run \u2014 no files written)")


if __name__ == "__main__":
    main()
