"""
Codegen: gidopensees schemas.json -> Pydantic v2 model stubs.

CLI:
    python -m tools.gidopensees_import.codegen \\
        --schemas tools/gidopensees_import/schemas.json \\
        --out src/otko/core/catalog/generated/
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

# ---------------------------------------------------------------------------
# Name helpers
# ---------------------------------------------------------------------------


def _to_snake(name: str) -> str:
    """Convert a GiD name to a valid Python snake_case identifier."""
    s = re.sub(r"[^a-zA-Z0-9]", "_", name)
    s = re.sub(r"_+", "_", s).strip("_")
    return s.lower() if s else "field"


def _to_class_name(name: str) -> str:
    """Convert a GiD entry name to PascalCaseSpec."""
    s = re.sub(r"[^a-zA-Z0-9]", "_", name)
    return "".join(p.capitalize() for p in s.split("_") if p) + "Spec"


def _to_module_name(name: str) -> str:
    """Convert a GiD entry name to a Python module filename stem."""
    s = re.sub(r"[^a-zA-Z0-9]", "_", name)
    s = re.sub(r"_+", "_", s).strip("_")
    return s.lower()


def _to_book_var(book_name: str) -> str:
    """Convert 'Standard_Uniaxial_Materials' -> 'StandardUniaxialMaterials'."""
    s = re.sub(r"[^a-zA-Z0-9]", "_", book_name)
    return "".join(p.capitalize() for p in s.split("_") if p)


def _dedupe(names: list[str]) -> list[str]:
    """Add _N suffix to duplicate identifiers, in order."""
    seen: dict[str, int] = {}
    result: list[str] = []
    for name in names:
        count = seen.get(name, 0)
        result.append(name if count == 0 else f"{name}_{count}")
        seen[name] = count + 1
    return result


# ---------------------------------------------------------------------------
# Type-inference helpers
# ---------------------------------------------------------------------------


def _scalar_type(default: str) -> tuple[str, str]:
    """Return (python_type_str, repr_default) for a SCALAR field."""
    if not default:
        return "str", '""'
    try:
        int(default)
        return "int", default
    except ValueError:
        pass
    try:
        float(default)
        return "float", default
    except ValueError:
        pass
    return "str", repr(default)


# ---------------------------------------------------------------------------
# Python-literal renderer (for model_config json_schema_extra)
# ---------------------------------------------------------------------------


def _py_lit(obj: Any, depth: int = 0) -> str:
    """Render a Python value as a valid Python literal expression."""
    pad = "    " * depth
    inner = "    " * (depth + 1)
    if isinstance(obj, dict):
        if not obj:
            return "{}"
        items = [f"{inner}{k!r}: {_py_lit(v, depth + 1)}" for k, v in obj.items()]
        return "{\n" + ",\n".join(items) + f",\n{pad}}}"
    if isinstance(obj, list):
        if not obj:
            return "[]"
        items = [f"{inner}{_py_lit(v, depth + 1)}" for v in obj]
        return "[\n" + ",\n".join(items) + f",\n{pad}]"
    return repr(obj)


# ---------------------------------------------------------------------------
# File header template
# ---------------------------------------------------------------------------

_HEADER = """\
# ruff: noqa
# AUTOGENERATED - do not edit.
# File: {module_name}.py
# Generated: {timestamp}
# Source: {source_file} -- Book: {book} / {entry_type}: {entry_name}
#
# Schema data from gidopensees (https://github.com/rclab-auth/gidopensees)
# Copyright (C) Reinforced Concrete Laboratory,
#               Aristotle University of Thessaloniki (AUTh)
"""

# ---------------------------------------------------------------------------
# Per-entry source generator
# ---------------------------------------------------------------------------


def _entry_source(
    entry: dict[str, Any],
    module_name: str,
    class_name: str,
    source_file: str,
    timestamp: str,
) -> str:
    lines: list[str] = []

    # Header block
    lines.append(
        _HEADER.format(
            module_name=module_name,
            timestamp=timestamp,
            source_file=source_file,
            book=entry["book"],
            entry_type=entry["source_type"],
            entry_name=entry["name"],
        )
    )

    # Visible fields
    visible = [f for f in entry["fields"] if f.get("state") != "HIDDEN"]

    # Detect which imports are needed
    has_literal = any(f["widget_type"] == "CB" and f.get("options") for f in visible)
    tuple_fields = [f for f in visible if f["widget_type"] == "TUPLE" and f.get("options")]

    # Inner row classes for TUPLE fields
    inner_classes: list[tuple[str, list[str]]] = []
    seen_inner: set[str] = set()
    for tf in tuple_fields:
        inner_cls = (
            "".join(p.capitalize() for p in _to_snake(tf["name"]).split("_") if p) + "Row"
        )
        if inner_cls not in seen_inner:
            seen_inner.add(inner_cls)
            rows: list[str] = [f"class {inner_cls}(BaseModel):"]
            for sub in tf.get("options") or []:
                rows.append(f"    {_to_snake(sub)}: float = 0.0")
            inner_classes.append((inner_cls, rows))

    # Imports
    lines.append("from __future__ import annotations\n\n")
    if has_literal:
        lines.append("from typing import Literal\n\n")
    lines.append("from pydantic import BaseModel, ConfigDict, Field\n\n\n")

    # Inner models
    for _, cls_rows in inner_classes:
        for row in cls_rows:
            lines.append(row + "\n")
        lines.append("\n\n")

    # Collect per-field dependency metadata for model_config
    deps: list[dict[str, Any]] = []
    for f in visible:
        if f.get("dependencies"):
            deps.append(
                {
                    "field": _to_snake(f["name"]),
                    "gid_name": f["name"],
                    "rules": f["dependencies"],
                }
            )

    # Build json_schema_extra dict
    extra: dict[str, Any] = {
        "x-gid-name": entry["name"],
        "x-book": entry["book"],
    }
    if entry.get("image"):
        extra["x-icon"] = entry["image"]
    if entry.get("tkwidgets"):
        extra["x-tkwidgets"] = entry["tkwidgets"]
    if deps:
        extra["dependencies"] = deps

    # Class header + model_config
    lines.append(f"class {class_name}(BaseModel):\n")
    extra_lit = _py_lit(extra, depth=2)
    lines.append(f"    model_config = ConfigDict(\n        json_schema_extra={extra_lit}\n    )\n\n")

    # Deduplicated Python attribute names
    py_names = _dedupe([_to_snake(f["name"]) for f in visible])

    for py_name, field in zip(py_names, visible, strict=True):
        wtype: str = field["widget_type"]
        options: list[str] = field.get("options") or []
        default: str = field.get("default") or ""
        help_text: str = field.get("help_text") or ""
        desc_arg = f", description={help_text!r}" if help_text else ""

        if wtype == "CB":
            if options:
                opts_str = ", ".join(repr(o) for o in options)
                py_type = f"Literal[{opts_str}]"
                py_def = repr(default) if default in options else repr(options[0])
            else:
                py_type, py_def = "str", repr(default)
            if desc_arg:
                lines.append(f"    {py_name}: {py_type} = Field({py_def}{desc_arg})\n")
            else:
                lines.append(f"    {py_name}: {py_type} = {py_def}\n")

        elif wtype == "UNITS":
            py_def = repr(default)
            if desc_arg:
                lines.append(
                    f"    {py_name}: str = Field({py_def}{desc_arg})"
                    "  # TODO: unit-aware type\n"
                )
            else:
                lines.append(f"    {py_name}: str = {py_def}  # TODO: unit-aware type\n")

        elif wtype == "MAT":
            py_def = repr(default)
            if desc_arg:
                lines.append(f"    {py_name}: str = Field({py_def}{desc_arg})\n")
            else:
                lines.append(f"    {py_name}: str = {py_def}\n")

        elif wtype == "SCALAR":
            py_type, py_def = _scalar_type(default)
            if desc_arg:
                lines.append(f"    {py_name}: {py_type} = Field({py_def}{desc_arg})\n")
            else:
                lines.append(f"    {py_name}: {py_type} = {py_def}\n")

        elif wtype == "TUPLE":
            inner_cls = (
                "".join(p.capitalize() for p in _to_snake(field["name"]).split("_") if p)
                + "Row"
            )
            lines.append(
                f"    {py_name}: list[{inner_cls}] = Field(default_factory=list{desc_arg})\n"
            )

        else:
            lines.append(f"    {py_name}: str = {default!r}\n")

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


# ---------------------------------------------------------------------------
# generated/__init__.py builder
# ---------------------------------------------------------------------------


def _generated_init_source(
    books: list[dict[str, Any]],
    source: str,
    module_prefix: str,
    timestamp: str,
) -> str:
    lines: list[str] = []
    lines.append("# ruff: noqa\n")
    lines.append(f"# AUTOGENERATED - do not edit.  Generated: {timestamp}\n")
    lines.append(
        "# Schema data from gidopensees.\n\n"
    )
    lines.append("from __future__ import annotations\n\n")
    lines.append("from typing import Union\n\n")

    # Collect (module_stem, class_name) per book
    book_entries: list[tuple[str, list[tuple[str, str]]]] = []
    for book in books:
        pairs: list[tuple[str, str]] = []
        for entry in book["entries"]:
            stem = _to_module_name(entry["name"])
            cls = _to_class_name(entry["name"])
            pairs.append((stem, cls))
        book_entries.append((book["name"], pairs))

    # Import all
    for _, pairs in book_entries:
        for stem, cls in pairs:
            lines.append(f"from {module_prefix}.{stem} import {cls}\n")
    lines.append("\n\n")

    # One Union per book
    for book_name, pairs in book_entries:
        var_name = _to_book_var(book_name)
        union_members = ", ".join(cls for _, cls in pairs)
        if len(pairs) == 1:
            lines.append(f"{var_name} = {pairs[0][1]}\n")
        else:
            lines.append(f"{var_name} = Union[{union_members}]\n")
    lines.append("\n")

    # __all__
    all_names = [_to_book_var(b["name"]) for b in books]
    all_cls = [cls for _, pairs in book_entries for _, cls in pairs]
    all_items = ", ".join(repr(n) for n in all_cls + all_names)
    lines.append(f"__all__ = [{all_items}]\n")

    return "".join(lines)


# ---------------------------------------------------------------------------
# core/catalog/__init__.py builder
# ---------------------------------------------------------------------------


def _catalog_init_source(
    mat_books: list[dict[str, Any]],
    timestamp: str,
) -> str:
    lines: list[str] = []
    lines.append("# ruff: noqa\n")
    lines.append(f"# AUTOGENERATED - do not edit.  Generated: {timestamp}\n")
    lines.append("# Flat CATALOG mapping gidopensees material name -> Pydantic Spec class.\n\n")
    lines.append("from __future__ import annotations\n\n")
    lines.append("from pydantic import BaseModel\n\n")

    base_pkg = "otko.core.catalog.generated"
    entries: list[tuple[str, str, str]] = []  # (gid_name, stem, class_name)
    for book in mat_books:
        for entry in book["entries"]:
            gid_name = entry["name"]
            stem = _to_module_name(gid_name)
            cls = _to_class_name(gid_name)
            entries.append((gid_name, stem, cls))

    for _gid_name, stem, cls in entries:
        lines.append(f"from {base_pkg}.{stem} import {cls}\n")
    lines.append("\n\n")

    lines.append("CATALOG: dict[str, type[BaseModel]] = {\n")
    for gid_name, _, cls in entries:
        lines.append(f"    {gid_name!r}: {cls},\n")
    lines.append("}\n\n")

    all_items = ", ".join(repr(cls) for _, _, cls in entries)
    lines.append(f'__all__ = ["CATALOG", {all_items}]\n')

    return "".join(lines)


# ---------------------------------------------------------------------------
# Orchestration
# ---------------------------------------------------------------------------


def run_codegen(schemas_path: Path, out_dir: Path) -> None:
    data: dict[str, Any] = json.loads(schemas_path.read_text(encoding="utf-8"))
    mat_books: list[dict[str, Any]] = data["mat_books"]
    cnd_books: list[dict[str, Any]] = data["cnd_books"]

    timestamp = datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    source_file = schemas_path.name

    # Ensure output directories
    cnd_dir = out_dir / "conditions"
    out_dir.mkdir(parents=True, exist_ok=True)
    cnd_dir.mkdir(parents=True, exist_ok=True)

    mat_count = 0
    cnd_count = 0

    # --- Material stubs ---
    for book in mat_books:
        for entry in book["entries"]:
            stem = _to_module_name(entry["name"])
            cls = _to_class_name(entry["name"])
            src = _entry_source(
                entry=entry,
                module_name=stem,
                class_name=cls,
                source_file=source_file,
                timestamp=timestamp,
            )
            (out_dir / f"{stem}.py").write_text(src, encoding="utf-8")
            mat_count += 1

    # --- Condition stubs ---
    for book in cnd_books:
        for entry in book["entries"]:
            stem = _to_module_name(entry["name"])
            cls = _to_class_name(entry["name"])
            src = _entry_source(
                entry=entry,
                module_name=stem,
                class_name=cls,
                source_file=source_file,
                timestamp=timestamp,
            )
            (cnd_dir / f"{stem}.py").write_text(src, encoding="utf-8")
            cnd_count += 1

    # --- generated/__init__.py ---
    mat_pkg = "otko.core.catalog.generated"
    gen_init = _generated_init_source(mat_books, "mat", mat_pkg, timestamp)
    (out_dir / "__init__.py").write_text(gen_init, encoding="utf-8")

    # --- generated/conditions/__init__.py ---
    cnd_pkg = "otko.core.catalog.generated.conditions"
    cnd_init = _generated_init_source(cnd_books, "cnd", cnd_pkg, timestamp)
    (cnd_dir / "__init__.py").write_text(cnd_init, encoding="utf-8")

    # --- core/catalog/__init__.py (parent of generated/) ---
    catalog_dir = out_dir.parent
    cat_init = _catalog_init_source(mat_books, timestamp)
    (catalog_dir / "__init__.py").write_text(cat_init, encoding="utf-8")

    # --- core/catalog/curated/__init__.py (empty placeholder) ---
    curated_dir = catalog_dir / "curated"
    curated_dir.mkdir(exist_ok=True)
    curated_init = curated_dir / "__init__.py"
    if not curated_init.exists():
        curated_init.write_text("", encoding="utf-8")

    print(f"Generated {mat_count} material stubs -> {out_dir}")
    print(f"Generated {cnd_count} condition stubs -> {cnd_dir}")
    print(f"Written -> {out_dir / '__init__.py'}")
    print(f"Written -> {cnd_dir / '__init__.py'}")
    print(f"Written -> {catalog_dir / '__init__.py'}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def _cli(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(
        description="Generate Pydantic v2 stubs from gidopensees schemas.json"
    )
    parser.add_argument(
        "--schemas",
        required=True,
        type=Path,
        help="Path to schemas.json (e.g. tools/gidopensees_import/schemas.json)",
    )
    parser.add_argument(
        "--out",
        required=True,
        type=Path,
        help="Output directory for generated stubs (e.g. src/otko/core/catalog/generated/)",
    )
    args = parser.parse_args(argv)

    schemas_path: Path = args.schemas
    out_dir: Path = args.out

    if not schemas_path.exists():
        print(f"error: schemas file not found: {schemas_path}", file=sys.stderr)
        sys.exit(1)

    run_codegen(schemas_path, out_dir)


if __name__ == "__main__":
    _cli()