from __future__ import annotations

from copy import deepcopy
from pathlib import Path
from typing import Any

from jsonschema import Draft202012Validator, FormatChecker

from .util import load_json, project_root, write_json


def canonical_schema_path() -> Path:
    return project_root() / "schema" / "page.schema.json"


def model_schema_path() -> Path:
    return project_root() / "schema" / "model-output.schema.json"


def load_canonical_schema() -> dict[str, Any]:
    return load_json(canonical_schema_path())


def _referenced_defs(node: Any) -> set[str]:
    names: set[str] = set()
    if isinstance(node, dict):
        reference = node.get("$ref")
        prefix = "#/$defs/"
        if isinstance(reference, str) and reference.startswith(prefix):
            names.add(reference.removeprefix(prefix))
        for value in node.values():
            names.update(_referenced_defs(value))
    elif isinstance(node, list):
        for value in node:
            names.update(_referenced_defs(value))
    return names


def _reachable_defs(root: dict[str, Any], definitions: dict[str, Any]) -> dict[str, Any]:
    """Keep only definitions reachable from the model-facing root schema.

    Besides making the schema smaller, this keeps canonical-only constraints out of
    provider structured-output validators. The canonical schema remains the final
    authority and is always applied after generation.
    """

    pending = sorted(_referenced_defs(root))
    reachable: dict[str, Any] = {}
    while pending:
        name = pending.pop(0)
        if name in reachable:
            continue
        definition = definitions[name]
        reachable[name] = deepcopy(definition)
        newly_referenced = _referenced_defs(definition) - reachable.keys() - set(pending)
        pending.extend(sorted(newly_referenced))
        pending.sort()
    return {name: reachable[name] for name in sorted(reachable)}


def _flatten_all_of(node: Any, definitions: dict[str, Any]) -> Any:
    """Flatten the simple object inheritance used by the canonical schema."""

    if isinstance(node, list):
        return [_flatten_all_of(value, definitions) for value in node]
    if not isinstance(node, dict):
        return node
    if "allOf" not in node:
        return {
            key: _flatten_all_of(value, definitions) for key, value in node.items()
        }

    merged: dict[str, Any] = {}
    merged_properties: dict[str, Any] = {}
    merged_required: list[str] = []
    for raw_branch in node["allOf"]:
        branch = raw_branch
        if isinstance(branch, dict) and set(branch) == {"$ref"}:
            prefix = "#/$defs/"
            reference = branch["$ref"]
            if reference.startswith(prefix):
                branch = definitions[reference.removeprefix(prefix)]
        branch = _flatten_all_of(branch, definitions)
        if not isinstance(branch, dict):
            continue
        for key, value in branch.items():
            if key == "properties":
                for property_name, property_schema in value.items():
                    # Canonical subtypes use true as a compact way to inherit the
                    # base property's constraint. Keep the inherited constraint.
                    if property_schema is True and property_name in merged_properties:
                        continue
                    merged_properties[property_name] = property_schema
            elif key == "required":
                for property_name in value:
                    if property_name not in merged_required:
                        merged_required.append(property_name)
            else:
                merged[key] = value
    for key, value in node.items():
        if key != "allOf":
            merged[key] = _flatten_all_of(value, definitions)
    if merged_properties:
        merged["properties"] = merged_properties
    if merged_required:
        merged["required"] = merged_required
    return merged


def _provider_compatible(node: Any) -> Any:
    """Translate equivalent constructs into the constrained-output subset.

    The Codex structured-output endpoint accepts ``anyOf`` but not ``oneOf``.
    Discriminated alternatives in this model are mutually exclusive through their
    ``type`` constants, so the substitution does not broaden accepted instances.
    """

    if isinstance(node, dict):
        translated: dict[str, Any] = {}
        for key, value in node.items():
            if key == "uniqueItems":
                # Enforced by the canonical validator after generation; unsupported
                # by the provider's constrained-output JSON Schema subset.
                continue
            translated["anyOf" if key == "oneOf" else key] = _provider_compatible(value)
        if "const" in translated and "type" not in translated:
            constant = translated["const"]
            if isinstance(constant, str):
                translated["type"] = "string"
            elif isinstance(constant, bool):
                translated["type"] = "boolean"
            elif isinstance(constant, int):
                translated["type"] = "integer"
            elif constant is None:
                translated["type"] = "null"
        return translated
    if isinstance(node, list):
        return [_provider_compatible(value) for value in node]
    return node


def export_model_schema(destination: Path | None = None) -> dict[str, Any]:
    canonical = load_canonical_schema()
    definitions = canonical["$defs"]
    flattened_definitions = {
        name: _flatten_all_of(definition, definitions)
        for name, definition in definitions.items()
    }
    draft = _flatten_all_of(deepcopy(definitions["modelDraft"]), definitions)
    exported = _provider_compatible({
        "$schema": canonical["$schema"],
        "$id": "https://crowley-diary.local/schema/model-output.schema.json",
        "title": "Crowley diary model page draft",
        **draft,
        "$defs": _reachable_defs(draft, flattened_definitions),
    })
    target = destination or model_schema_path()
    write_json(target, exported)
    return exported


def schema_errors(instance: Any, schema: dict[str, Any]) -> list[str]:
    validator = Draft202012Validator(schema, format_checker=FormatChecker())
    rendered: list[str] = []
    for error in sorted(validator.iter_errors(instance), key=lambda item: list(item.path)):
        path = "/" + "/".join(str(part) for part in error.absolute_path)
        rendered.append(f"{path or '/'}: {error.message}")
    return rendered


def check_schemas() -> None:
    canonical = load_canonical_schema()
    Draft202012Validator.check_schema(canonical)
    exported = export_model_schema()
    Draft202012Validator.check_schema(exported)
