from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

from scienith_openapi_contract_tooling.openapi.load import load_openapi_spec


def _split_ref(ref: str) -> tuple[str, str]:
    if "#" in ref:
        path_part, frag = ref.split("#", 1)
        return path_part, frag
    return ref, ""


def _json_pointer_get(doc: Any, pointer: str) -> Any:
    # Accept both "/a/b" and "a/b"
    if pointer.startswith("/"):
        pointer = pointer[1:]
    if pointer == "":
        return doc

    cur: Any = doc
    for raw_token in pointer.split("/"):
        token = raw_token.replace("~1", "/").replace("~0", "~")
        if isinstance(cur, list):
            idx = int(token)
            cur = cur[idx]
        elif isinstance(cur, dict):
            cur = cur[token]
        else:
            raise KeyError(f"JSON Pointer 不可达：{raw_token}")
    return cur


def _canonical_ref(base_dir: Path, ref: str) -> str:
    path_part, frag = _split_ref(ref)
    if path_part == "":
        # internal ref
        return f"#{frag}"
    p = (base_dir / path_part).resolve()
    return f"{p}#{frag}"


def _load_ref_target(*, base_dir: Path, ref: str) -> tuple[Any, Path]:
    path_part, frag = _split_ref(ref)
    if path_part == "":
        raise ValueError("内部 $ref 不能通过 _load_ref_target 加载。")

    target_path = (base_dir / path_part).resolve()
    doc = load_openapi_spec(target_path)
    if frag.startswith("/"):
        value = _json_pointer_get(doc, frag)
    elif frag == "":
        value = doc
    else:
        # allow "#/..." only
        raise ValueError(f"不支持的 $ref fragment：#{frag}")
    return value, target_path.parent


@dataclass
class ResolvedSchema:
    name: str
    schema: dict[str, Any]
    base_dir: Path
    source_ref: str | None = None


def resolve_components_schemas(*, root_path: Path) -> dict[str, ResolvedSchema]:
    root_path = root_path.resolve()
    root = load_openapi_spec(root_path)
    schemas = (((root.get("components") or {}).get("schemas")) or {})
    if not isinstance(schemas, dict) or not schemas:
        raise ValueError("未找到 components/schemas。")

    resolved: dict[str, ResolvedSchema] = {}
    for name, schema_or_ref in schemas.items():
        if isinstance(schema_or_ref, dict) and "$ref" in schema_or_ref:
            ref_str = str(schema_or_ref["$ref"])
            source_ref = _canonical_ref(root_path.parent, ref_str)
            target, base_dir = _load_ref_target(base_dir=root_path.parent, ref=ref_str)
            if not isinstance(target, dict):
                raise ValueError(f"schema $ref 目标不是对象：{name} -> {schema_or_ref['$ref']}")
            resolved[str(name)] = ResolvedSchema(
                name=str(name),
                schema=target,
                base_dir=base_dir,
                source_ref=source_ref,
            )
        elif isinstance(schema_or_ref, dict):
            resolved[str(name)] = ResolvedSchema(name=str(name), schema=schema_or_ref, base_dir=root_path.parent)
        else:
            raise ValueError(f"schema 定义非法：{name}")
    return resolved


def normalize_schema_refs(
    *,
    resolved: dict[str, ResolvedSchema],
    root_dir: Path,
) -> dict[str, dict[str, Any]]:
    """
    将 components/schemas 内部及其子树中的“文件 $ref”尽可能改写为
    `#/components/schemas/<Name>`，以便后续 UML 逻辑做稳定的“按 schema 名”关联。
    """
    root_dir = root_dir.resolve()

    file_ref_to_name: dict[str, str] = {}
    for name, rs in resolved.items():
        if rs.source_ref:
            file_ref_to_name[rs.source_ref] = name

    def rewrite(node: Any, *, base_dir: Path) -> Any:
        if isinstance(node, dict):
            if "$ref" in node and isinstance(node["$ref"], str):
                ref = node["$ref"]
                path_part, frag = _split_ref(ref)
                if path_part != "":
                    canonical = _canonical_ref(base_dir, ref)
                    mapped = file_ref_to_name.get(canonical)
                    if mapped:
                        return {"$ref": f"#/components/schemas/{mapped}"}
                return node

            return {k: rewrite(v, base_dir=base_dir) for k, v in node.items()}
        if isinstance(node, list):
            return [rewrite(v, base_dir=base_dir) for v in node]
        return node

    normalized: dict[str, dict[str, Any]] = {}
    for name, rs in resolved.items():
        normalized[name] = rewrite(rs.schema, base_dir=rs.base_dir)
    return normalized
