"""Load contracts/schema/ and lower it into a small IR the generators emit from.

There is deliberately only one of these. gen_cpp.py, gen_ts.py and gen_openrpc.py all
consume the same IR, so the C++ structs, the TypeScript types and the human-readable
API document cannot disagree about what the contract says.

The IR covers exactly the JSON Schema subset the contract is allowed to use. Anything
outside it raises SchemaError at generation time rather than producing subtly wrong code:
a contract that cannot be generated from is a contract bug, and it should stop the build.

Supported subset
----------------
  type: object / string / integer / number / boolean / array, and ["X", "null"]
  object with `properties`      -> struct
  object with `additionalProperties: <schema>` -> map<string, T>
  string with `enum`            -> enum
  integer with `enum` + x-enum  -> named integer enum
  array with `items`            -> vector<T>
  $ref to a types/*.schema.json -> named type
  oneOf: [{$ref}, {type: null}] -> nullable named type
  {}                            -> opaque JSON

Not supported, on purpose: allOf, anyOf, general oneOf, patternProperties, tuple items,
recursive types. If the contract needs one of these, extend this file in the same PR.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from pathlib import Path

SCHEMA_ROOT = Path(__file__).resolve().parent.parent / "schema"
ID_PREFIX = "https://velox.dev/schema/"


class SchemaError(Exception):
    """A schema the generators refuse to guess about."""


# --------------------------------------------------------------------------- IR


# Value constraints the generators enforce at runtime. The wire is never trusted, so a
# `maximum` in the schema has to be a check in the generated code, not just documentation.
CONSTRAINT_KEYS = ("minimum", "maximum", "minLength", "maxLength", "pattern",
                   "minItems", "maxItems")


@dataclass(frozen=True)
class TypeRef:
    """A reference to a type from a field. `kind` drives every emitter's switch."""

    kind: str  # named | string | integer | number | boolean | array | map | json
    name: str | None = None          # kind == "named"
    inner: "TypeRef | None" = None   # kind in ("array", "map")
    constraints: tuple[tuple[str, object], ...] = ()

    @property
    def limits(self) -> dict[str, object]:
        return dict(self.constraints)


@dataclass
class Field:
    name: str
    type: TypeRef
    required: bool
    nullable: bool
    doc: str = ""

    @property
    def optional(self) -> bool:
        """Whether the emitted field needs an optional/undefined-able representation."""
        return (not self.required) or self.nullable


@dataclass
class EnumValue:
    name: str      # identifier form, e.g. RetryWait
    wire: object   # the value on the wire: "retry_wait" or -32010
    doc: str = ""


@dataclass
class TypeDef:
    name: str
    kind: str  # struct | string_enum | int_enum | map_alias
    doc: str = ""
    fields: list[Field] = field(default_factory=list)
    values: list[EnumValue] = field(default_factory=list)
    alias: TypeRef | None = None
    source: str = ""  # relative path, for the "do not edit" banner


@dataclass
class Method:
    name: str                 # download.addBatch
    doc: str
    params: TypeRef
    result: TypeRef
    privileged: bool
    transports: list[str]
    deadline_ms: int
    errors: list[int] = field(default_factory=list)
    ws_restrictions: list[str] = field(default_factory=list)


@dataclass
class Event:
    name: str                 # event.task.progress
    doc: str
    params: TypeRef
    max_rate_hz: float | None = None


@dataclass
class Contract:
    version: str
    types: list[TypeDef]
    methods: list[Method]
    events: list[Event]
    error_codes: list[EnumValue]


# ------------------------------------------------------------------ name helpers


def pascal(text: str) -> str:
    """download.addBatch -> DownloadAddBatch ; retry_wait -> RetryWait."""
    parts = re.split(r"[.\-_ ]+", text)
    out = []
    for part in parts:
        if not part:
            continue
        out.append(part[0].upper() + part[1:])
    return "".join(out)


def enum_ident(value: str) -> str:
    ident = pascal(value)
    if not ident:
        raise SchemaError(f"cannot derive an identifier from enum value {value!r}")
    if ident[0].isdigit():
        ident = "V" + ident
    return ident


# ---------------------------------------------------------------------- loading


class Loader:
    def __init__(self, root: Path = SCHEMA_ROOT):
        self.root = root
        self.by_id: dict[str, dict] = {}
        self.source_of: dict[str, str] = {}
        for path in sorted(root.rglob("*.schema.json")):
            with path.open() as fh:
                doc = json.load(fh)
            sid = doc.get("$id")
            if not sid:
                raise SchemaError(f"{path} has no $id")
            if sid in self.by_id:
                raise SchemaError(f"duplicate $id {sid} in {path}")
            self.by_id[sid] = doc
            self.source_of[sid] = str(path.relative_to(root.parent))

        self.types: dict[str, TypeDef] = {}
        self._order: list[str] = []

    # -- ref handling -------------------------------------------------------

    def _resolve_ref(self, ref: str) -> dict:
        if ref.startswith("#"):
            raise SchemaError(f"local $ref {ref} is not supported outside envelope.schema.json")
        if ref not in self.by_id:
            raise SchemaError(f"unknown $ref {ref}")
        return self.by_id[ref]

    @staticmethod
    def _split_nullable(node: dict) -> tuple[dict, bool]:
        """Normalise the two ways the contract spells 'or null'."""
        if "oneOf" in node:
            branches = node["oneOf"]
            non_null = [b for b in branches if b.get("type") != "null"]
            nulls = [b for b in branches if b.get("type") == "null"]
            if len(branches) != 2 or len(non_null) != 1 or len(nulls) != 1:
                raise SchemaError(
                    "oneOf is only supported as [<schema>, {type: null}]; got "
                    + json.dumps(branches)[:200]
                )
            merged = dict(non_null[0])
            for key in ("description",):
                if key in node and key not in merged:
                    merged[key] = node[key]
            return merged, True

        t = node.get("type")
        if isinstance(t, list):
            non_null = [x for x in t if x != "null"]
            if len(non_null) != 1:
                raise SchemaError(f"union type {t} is only supported as [X, 'null']")
            node = dict(node)
            node["type"] = non_null[0]
            # An enum listing null alongside its values means the same thing.
            if "enum" in node:
                node["enum"] = [v for v in node["enum"] if v is not None]
            return node, True

        if "enum" in node and None in node["enum"]:
            node = dict(node)
            node["enum"] = [v for v in node["enum"] if v is not None]
            return node, True

        return node, False

    # -- lowering -----------------------------------------------------------

    def type_ref(self, node: dict, hint: str) -> tuple[TypeRef, bool]:
        """Lower a schema node to a TypeRef. `hint` names any struct we must synthesise."""
        node, nullable = self._split_nullable(node)
        limits = tuple((k, node[k]) for k in CONSTRAINT_KEYS if k in node)

        if "$ref" in node:
            target = self._resolve_ref(node["$ref"])
            return TypeRef("named", name=self.named_type(node["$ref"], target)), nullable

        if not node or node.keys() <= {"description"}:
            return TypeRef("json"), True

        t = node.get("type")
        if t == "string":
            if "enum" in node:
                return TypeRef("named", name=self._synth_string_enum(hint, node)), nullable
            return TypeRef("string", constraints=limits), nullable
        if t == "integer":
            return TypeRef("integer", constraints=limits), nullable
        if t == "number":
            return TypeRef("number", constraints=limits), nullable
        if t == "boolean":
            return TypeRef("boolean", constraints=limits), nullable
        if t == "array":
            items = node.get("items")
            if items is None:
                raise SchemaError(f"array without items at {hint}")
            inner, _ = self.type_ref(items, hint + "Item")
            return TypeRef("array", inner=inner, constraints=limits), nullable
        if t == "object":
            if "properties" in node:
                return TypeRef("named", name=self._synth_struct(hint, node)), nullable
            ap = node.get("additionalProperties")
            if isinstance(ap, dict):
                inner, _ = self.type_ref(ap, hint + "Value")
                return TypeRef("map", inner=inner), nullable
            return TypeRef("json"), nullable
        if "const" in node:
            return TypeRef("string", constraints=limits), nullable

        raise SchemaError(f"unsupported schema node at {hint}: {json.dumps(node)[:200]}")

    def _register(self, td: TypeDef) -> str:
        existing = self.types.get(td.name)
        if existing is not None:
            if existing.kind != td.kind:
                raise SchemaError(f"type name collision on {td.name}")
            return td.name
        self.types[td.name] = td
        self._order.append(td.name)
        return td.name

    def _synth_string_enum(self, name: str, node: dict) -> str:
        values = [
            EnumValue(name=enum_ident(v), wire=v)
            for v in node["enum"]
            if v is not None
        ]
        return self._register(TypeDef(name=name, kind="string_enum",
                                      doc=node.get("description", ""), values=values))

    def _synth_struct(self, name: str, node: dict) -> str:
        if node.get("additionalProperties", False) is not False:
            raise SchemaError(
                f"{name}: object schemas must set additionalProperties:false — the daemon is "
                "not allowed to trust unknown fields on the wire"
            )
        required = set(node.get("required", []))
        fields: list[Field] = []
        for prop, sub in node.get("properties", {}).items():
            ref, nullable = self.type_ref(sub, name + pascal(prop))
            fields.append(Field(name=prop, type=ref, required=prop in required,
                                nullable=nullable, doc=sub.get("description", "")))
        return self._register(TypeDef(name=name, kind="struct",
                                      doc=node.get("description", ""), fields=fields))

    def named_type(self, sid: str, doc: dict) -> str:
        """Lower a top-level types/*.schema.json into a TypeDef and return its name."""
        name = doc.get("title")
        if not name:
            raise SchemaError(f"{sid} has no title")
        if name in self.types:
            return name

        source = self.source_of.get(sid, "")
        node, _ = self._split_nullable(doc)
        t = node.get("type")

        if t == "string" and "enum" in node:
            # Placeholder first: enums cannot recurse, but registering early keeps the
            # ordering stable and mirrors the struct path below.
            values = [EnumValue(name=enum_ident(v), wire=v) for v in node["enum"] if v is not None]
            td = TypeDef(name, "string_enum", node.get("description", ""), values=values, source=source)
            return self._register(td)

        if t == "integer" and "x-enum" in node:
            values = [EnumValue(name=e["name"], wire=e["value"], doc=e.get("doc", ""))
                      for e in node["x-enum"]]
            td = TypeDef(name, "int_enum", node.get("description", ""), values=values, source=source)
            return self._register(td)

        if t == "object" and "properties" not in node:
            ap = node.get("additionalProperties")
            if not isinstance(ap, dict):
                raise SchemaError(f"{sid}: object type with neither properties nor a typed additionalProperties")
            inner, _ = self.type_ref(ap, name + "Value")
            td = TypeDef(name, "map_alias", node.get("description", ""),
                         alias=TypeRef("map", inner=inner), source=source)
            return self._register(td)

        if t == "object":
            # Reserve the name before descending so a nested synth cannot steal it.
            self.types[name] = TypeDef(name, "struct", node.get("description", ""), source=source)
            self._order.append(name)
            required = set(node.get("required", []))
            if node.get("additionalProperties", False) is not False:
                raise SchemaError(f"{sid}: object schemas must set additionalProperties:false")
            fields = []
            for prop, sub in node.get("properties", {}).items():
                ref, nullable = self.type_ref(sub, name + pascal(prop))
                fields.append(Field(prop, ref, prop in required, nullable, sub.get("description", "")))
            self.types[name].fields = fields
            return name

        raise SchemaError(f"{sid}: unsupported top-level type {t!r}")


# ---------------------------------------------------------------------- driver


def load() -> Contract:
    loader = Loader()
    root = loader.root.parent

    version = (root / "VERSION").read_text().strip()

    # Named types first, so their names win over any synthesised ones.
    for sid in sorted(loader.by_id):
        if "/types/" in sid:
            loader.named_type(sid, loader.by_id[sid])

    methods: list[Method] = []
    for sid in sorted(loader.by_id):
        if "/methods/" not in sid:
            continue
        doc = loader.by_id[sid]
        name = doc["title"]
        base = pascal(name)
        props = doc.get("properties", {})
        for half in ("params", "result"):
            if half not in props:
                raise SchemaError(f"{sid}: method schema must define both params and result")
        params, _ = loader.type_ref(props["params"], base + "Params")
        result, _ = loader.type_ref(props["result"], base + "Result")
        transports = doc.get("x-transports")
        if not transports:
            raise SchemaError(f"{sid}: x-transports is required")
        if "x-privileged" not in doc:
            raise SchemaError(f"{sid}: x-privileged is required")
        methods.append(Method(
            name=name, doc=doc.get("description", ""), params=params, result=result,
            privileged=bool(doc["x-privileged"]), transports=list(transports),
            deadline_ms=int(doc.get("x-deadlineMs", 5000)),
            errors=list(doc.get("x-errors", [])),
            ws_restrictions=list(doc.get("x-wsRestrictions", [])),
        ))

    events: list[Event] = []
    for sid in sorted(loader.by_id):
        if "/events/" not in sid:
            continue
        doc = loader.by_id[sid]
        name = doc["title"]
        ident = pascal(name[len("event."):] if name.startswith("event.") else name) + "Event"
        params, _ = loader.type_ref(doc["properties"]["params"], ident)
        events.append(Event(name=name, doc=doc.get("description", ""), params=params,
                            max_rate_hz=doc.get("x-maxRateHz")))

    error_codes = loader.types["ErrorCode"].values
    ordered = [loader.types[n] for n in loader._order]
    return Contract(version=version, types=ordered, methods=methods, events=events,
                    error_codes=error_codes)


def dependencies(td: TypeDef) -> set[str]:
    """Named types `td` mentions directly."""
    out: set[str] = set()

    def walk(ref: TypeRef | None) -> None:
        if ref is None:
            return
        if ref.kind == "named" and ref.name:
            out.add(ref.name)
        walk(ref.inner)

    for f in td.fields:
        walk(f.type)
    walk(td.alias)
    return out


def topo_sorted(types: list[TypeDef]) -> list[TypeDef]:
    """Definition order for languages that need a type declared before it is used.

    The contract forbids recursive types, so a cycle here means a schema bug and is
    raised rather than broken arbitrarily.
    """
    by_name = {t.name: t for t in types}
    state: dict[str, int] = {}
    order: list[TypeDef] = []

    def visit(name: str, trail: list[str]) -> None:
        mark = state.get(name, 0)
        if mark == 2:
            return
        if mark == 1:
            raise SchemaError("recursive type: " + " -> ".join(trail + [name]))
        state[name] = 1
        for dep in sorted(dependencies(by_name[name])):
            if dep in by_name:
                visit(dep, trail + [name])
        state[name] = 2
        order.append(by_name[name])

    for t in types:
        visit(t.name, [])
    return order


if __name__ == "__main__":
    c = load()
    print(f"contract {c.version}: {len(c.types)} types, {len(c.methods)} methods, {len(c.events)} events")
    for t in c.types:
        detail = f"{len(t.fields)} fields" if t.kind == "struct" else f"{len(t.values)} values"
        print(f"  {t.kind:12} {t.name:34} {detail}")