"""Light, safe normalization of real-world JSON Schemas into deg-compilable ones.

Used for external schema sources (JSONSchemaBench for datagen, ExtractBench for eval). The
result is a self-contained schema that :func:`deg.schema.compile_schema` accepts, or
``None`` with a reason. Only lossless-or-annotation edits and *dropping whole fields* are
done; a field's meaning is never rewritten:

- local ``$ref`` is inlined (recursive refs drop the referring field), ``$defs`` removed;
- annotation-like and unsupported validation keywords are dropped (``examples``, vendor
  ``x-*``/custom keys, ``minProperties``, ``dependencies``, ``contentEncoding``, ...);
- ``format`` values deg does not parse are dropped (the field becomes a plain span);
- ``allOf`` with non-conflicting branches is merged; nullable objects/arrays lose ``null``;
- a plain ``anyOf`` of string and number types becomes a string span (nullable if listed);
- a missing ``type`` is inferred from ``properties``/``items``/``enum`` when unambiguous;
- fields that still cannot compile (``const``, free-form maps, tuple arrays, multi-type
  unions, non-enum ``anyOf``) are dropped, from ``required`` too;
- the tree is capped at ``max_leaves`` leaves and ``max_depth`` nesting (later properties
  are dropped first).

The compiler is the judge: after the rewrite, every field named in a compile error is
dropped and the schema recompiled until it compiles or nothing fillable is left.
"""

from __future__ import annotations

import copy
from typing import Any

from deg.schema import ANNOTATIONS, TYPE_KEYWORDS, SchemaError, compile_schema

FORMATS = {"date", "date-time", "email", "uri", "uuid"}
KEEP_ANNOTATIONS = {"title", "description"}
_VALIDATION = frozenset().union(*TYPE_KEYWORDS.values()) | {
    "enum", "const", "anyOf", "oneOf", "allOf", "not", "$ref"}


class Skip(ValueError):
    """The schema (or a subtree) cannot be normalized; ``reason`` is machine-readable."""

    def __init__(self, reason: str) -> None:
        super().__init__(reason)
        self.reason = reason


def _resolve_ref(root: Any, ref: str) -> Any:
    if not isinstance(ref, str) or not (ref == "#" or ref.startswith("#/")):
        raise Skip("nonlocal_ref")
    target = root
    for raw in ref[2:].split("/") if ref != "#" else []:
        part = raw.replace("~1", "/").replace("~0", "~")
        if isinstance(target, dict) and part in target:
            target = target[part]
        elif isinstance(target, list) and part.isdigit() and int(part) < len(target):
            target = target[int(part)]
        else:
            raise Skip("unresolvable_ref")
    if not isinstance(target, dict):
        raise Skip("bad_ref")
    return target


class _Normalizer:
    def __init__(self, root: dict[str, Any], max_depth: int) -> None:
        self.root = root
        self.max_depth = max_depth

    def inline(self, s: Any, refs: tuple[str, ...]) -> tuple[dict[str, Any], tuple[str, ...]]:
        while isinstance(s, dict) and "$ref" in s:
            ref = s["$ref"]
            if ref in refs:
                raise Skip("recursive_ref")
            target = _resolve_ref(self.root, ref)
            s = {**target, **{k: v for k, v in s.items() if k != "$ref"}}
            refs = (*refs, ref)
        if not isinstance(s, dict):
            raise Skip("boolean_schema")
        return s, refs

    def merge_all_of(self, s: dict[str, Any], refs: tuple[str, ...]) -> dict[str, Any]:
        out = {k: v for k, v in s.items() if k != "allOf"}
        for branch in s["allOf"] if isinstance(s["allOf"], list) else []:
            try:
                b, _ = self.inline(branch, refs)
            except Skip as exc:
                if exc.reason == "nonlocal_ref":  # an external base type: keep the rest
                    continue
                raise
            if "allOf" in b:
                b = self.merge_all_of(b, refs)
            for k, v in b.items():
                if k == "properties" and isinstance(v, dict):
                    out["properties"] = {**out.get("properties", {}), **v}
                elif k == "required" and isinstance(v, list):
                    out["required"] = list(dict.fromkeys([*out.get("required", []), *v]))
                elif k in out and out[k] != v and k not in ANNOTATIONS:
                    raise Skip("allOf_conflict")
                else:
                    out.setdefault(k, v)
        return out

    def node(self, raw: Any, refs: tuple[str, ...], depth: int) -> dict[str, Any]:
        s, refs = self.inline(raw, refs)
        if "allOf" in s:
            s = self.merge_all_of(s, refs)
        if "const" in s:
            raise Skip("const")
        for key in ("anyOf", "oneOf"):
            if key in s:
                return self.union(s, key, refs, depth)
        t = s.get("type")
        if t is None:
            if "properties" in s:
                t = "object"
            elif "items" in s:
                t = "array"
            elif isinstance(s.get("enum"), list) and s["enum"]:
                kinds = {_kind(v) for v in s["enum"] if v is not None}
                if len(kinds) != 1:
                    raise Skip("untyped_enum")
                t = kinds.pop()
            else:
                raise Skip("untyped")
        nullable = False
        if isinstance(t, list):
            non_null = [x for x in t if x != "null"]
            if len(non_null) != 1:
                raise Skip("multi_type")
            nullable = "null" in t
            t = non_null[0]
        if t not in TYPE_KEYWORDS:
            raise Skip(f"type_{t}")
        out: dict[str, Any] = {k: v for k, v in s.items()
                               if k in KEEP_ANNOTATIONS and isinstance(v, str)}
        if "enum" in s:
            if not isinstance(s["enum"], list):
                raise Skip("bad_enum")
            values = list(dict.fromkeys(_freeze(v) for v in s["enum"]))
            values = [_thaw(v) for v in values]
            if len([v for v in values if v is not None]) < 2 or len(values) > 255:
                raise Skip("enum_size")
            out["enum"] = values
        for k in TYPE_KEYWORDS[t]:
            if k in s and k not in ("properties", "items", "required", "additionalProperties"):
                if "enum" in out and t in ("string", "integer", "number"):
                    continue  # the enum already pins the value; bounds/lengths are redundant
                out[k] = s[k]
        if t == "string" and out.get("format") not in FORMATS:
            out.pop("format", None)
        if t == "object":
            if depth >= self.max_depth:
                raise Skip("too_deep")
            props = s.get("properties")
            if not isinstance(props, dict) or not props:
                raise Skip("free_form_object")
            req = s.get("required") if isinstance(s.get("required"), list) else []
            kept: dict[str, Any] = {}
            for name, sub in props.items():
                try:
                    kept[name] = self.node(sub, refs, depth + 1)
                except Skip:
                    continue
            if not kept:
                raise Skip("no_fillable_properties")
            out["properties"] = kept
            required = [r for r in req if r in kept]
            if required:
                out["required"] = required
            out["additionalProperties"] = False
            out["type"] = "object"  # nullable objects are not supported: drop null
            return out
        if t == "array":
            items = s.get("items")
            if not isinstance(items, dict):
                raise Skip("tuple_or_untyped_array")
            out["items"] = self.node(items, refs, depth + 1)
            if out["items"].get("type") == "array":
                raise Skip("nested_array")
            out["type"] = "array"
            return out
        if nullable and "enum" in out and None not in out["enum"]:
            out["enum"] = [*out["enum"], None]
        nullable = nullable or None in out.get("enum", [])
        out["type"] = [t, "null"] if nullable else t
        return out

    def union(self, s: dict[str, Any], key: str, refs: tuple[str, ...],
              depth: int) -> dict[str, Any]:
        branches = s[key] if isinstance(s[key], list) else []
        resolved = [self.inline(b, refs)[0] for b in branches]
        nulls = [b for b in resolved if b.get("type") == "null"]
        others = [b for b in resolved if b.get("type") != "null"]
        ann = {k: v for k, v in s.items() if k in KEEP_ANNOTATIONS and isinstance(v, str)}
        if len(others) == 1:  # nullable form (or a 1-branch union)
            inner = self.node({**others[0], **ann}, refs, depth)
            if nulls and inner.get("type") not in ("object", "array"):
                if "enum" in inner and None not in inner["enum"]:
                    inner["enum"] = [*inner["enum"], None]
                if isinstance(inner.get("type"), str):
                    inner["type"] = [inner["type"], "null"]
            return inner
        types = {b.get("type") if isinstance(b.get("type"), str) else "?" for b in resolved}
        plain = not any(set(b) & _VALIDATION for b in resolved)  # only type + annotations
        if plain and "string" in types and types <= {"string", "integer", "number", "null"}:
            # "a string or a number": the value is copied from text either way -> a span
            return self.node({**ann, "type": ["string", "null"] if nulls else "string"},
                             refs, depth)
        values: list[Any] = []
        for b in resolved:
            if "const" in b:
                values.append(b["const"])
            elif isinstance(b.get("enum"), list):
                values.extend(b["enum"])
            elif b.get("type") == "null":
                values.append(None)
            else:
                raise Skip("union")
        return self.node({**ann, "enum": values}, refs, depth)


def _kind(v: Any) -> str:
    if isinstance(v, bool):
        return "boolean"
    if isinstance(v, int):
        return "integer"
    if isinstance(v, float):
        return "number"
    if isinstance(v, str):
        return "string"
    return "object"


def _freeze(v: Any) -> Any:
    return ("__json__", repr(v)) if isinstance(v, (dict, list)) else (type(v).__name__, v)


def _thaw(v: Any) -> Any:
    if v[0] == "__json__":
        raise Skip("structured_enum")
    return v[1]


def count_leaves(s: dict[str, Any]) -> int:
    if s.get("type") == "object":
        return sum(count_leaves(p) for p in s["properties"].values())
    if s.get("type") == "array":
        return count_leaves(s["items"])
    return 1


def cap_leaves(s: dict[str, Any], budget: int) -> int:
    """Keep properties in order while they fit ``budget`` leaves; returns leaves kept."""
    if s.get("type") == "array":
        return cap_leaves(s["items"], budget)
    if s.get("type") != "object":
        return 1
    used = 0
    kept: dict[str, Any] = {}
    for name, sub in s["properties"].items():
        if used >= budget:
            break
        n = count_leaves(sub)
        if n > budget - used:
            if sub.get("type") in ("object", "array") and budget - used >= 2:
                n = cap_leaves(sub, budget - used)
            else:
                continue
        kept[name] = sub
        used += n
    s["properties"] = kept
    if "required" in s:
        s["required"] = [r for r in s["required"] if r in kept]
        if not s["required"]:
            del s["required"]
    return used


def _drop_at(schema: dict[str, Any], loc: list[Any]) -> bool:
    """Drop the deepest property named in an error ``loc``; False if none can be dropped."""
    if loc and loc[0] == "schema":  # compiler locs are rooted at "schema"
        loc = loc[1:]
    cuts = [i for i in range(len(loc) - 1) if loc[i] == "properties"]
    for i in reversed(cuts):
        parent: Any = schema
        try:
            for part in loc[:i]:
                parent = parent[part]
        except (KeyError, IndexError, TypeError):
            continue
        name = loc[i + 1]
        props = parent.get("properties") if isinstance(parent, dict) else None
        if isinstance(props, dict) and name in props and len(props) > 1:
            del props[name]
            if name in parent.get("required", []):
                parent["required"] = [r for r in parent["required"] if r != name]
                if not parent["required"]:
                    del parent["required"]
            return True
    return False


def normalize_schema(schema: Any, max_leaves: int = 30, max_depth: int = 5,
                     min_leaves: int = 2) -> tuple[dict[str, Any] | None, str]:
    """(deg-compilable schema, "ok") or (None, skip reason)."""
    if not isinstance(schema, dict):
        return None, "not_object"
    root = copy.deepcopy(schema)
    try:
        out = _Normalizer(root, max_depth).node(root, ("#",), 0)
    except Skip as exc:
        return None, exc.reason
    except RecursionError:
        return None, "recursion"
    if out.get("type") != "object":
        return None, "root_not_object"
    cap_leaves(out, max_leaves)
    for _ in range(64):
        try:
            compile_schema(out)
            break
        except SchemaError as exc:
            if not any(_drop_at(out, e["loc"]) for e in exc.errors[:1]):
                return None, "compile:" + exc.errors[0]["msg"][:60]
    else:
        return None, "compile_loop"
    if count_leaves(out) < min_leaves:
        return None, "too_few_leaves"
    return out, "ok"