#!/usr/bin/env python3

"""Internal utility functions

![status](../../badges/status-utils.svg)
"""

from __future__ import annotations

import importlib.util
import re
from collections.abc import Callable
from typing import Any

from pydantic import BaseModel

_CROSSREF_RE = re.compile(r"\[([^\]\r\n]+)\]\[([^\]\r\n]+)\]")


def make_crossref_resolver(
    docs_base_url: str = "https://psandhaas.github.io/evaLLM/reference",
) -> Callable[[dict[str, Any], type[BaseModel] | None], None]:  # pragma: no cover
    """Creates a function that resolves mkdocstrings-style cross-references in JSON
    schema descriptions.

    The returned function can be used as a callable for
    [`pydantic.BaseModel.json_schema_extra`](https://docs.pydantic.dev/latest/api/config/#pydantic.config.ConfigDict.json_schema_extra)
    to resolve cross-reference-paths in the generated JSON schema to URLs.

    Args:
        docs_base_url: The base URL of the documentation where the cross-references should
            point to. The resolver will append the appropriate path and anchor to this base
            URL based on the target of the cross-reference.

    Returns:
        Callable[dict[str, Any], type[BaseModel] | None]: A function that can be used
            as a `json_schema_extra` callable to resolve cross-references in JSON
            schema descriptions.
    """
    docs_base_url = docs_base_url.rstrip("/")

    def _iter_model_subclasses(cls: type[BaseModel]) -> set[type[BaseModel]]:
        found: set[type[BaseModel]] = set()
        stack = [cls]

        while stack:
            current = stack.pop()
            for subcls in current.__subclasses__():
                if subcls not in found:
                    found.add(subcls)
                    if issubclass(subcls, BaseModel):
                        stack.append(subcls)

        return found

    def _shared_model_root(model: type[BaseModel]) -> type[BaseModel]:
        candidates = [
            cls
            for cls in model.__mro__
            if issubclass(cls, BaseModel) and cls is not BaseModel
        ]
        return candidates[-1] if candidates else model

    def _build_def_registry(model: type[BaseModel]) -> dict[str, str]:
        root = _shared_model_root(model)
        classes = {root, *_iter_model_subclasses(root)}

        registry: dict[str, str] = {}
        for cls in classes:
            fqdn = f"{cls.__module__}.{cls.__name__}"
            registry[cls.__name__] = fqdn

            title = getattr(cls, "model_config", {}).get("title")
            if isinstance(title, str) and title:
                registry[title] = fqdn

        return registry

    def _count_leading_dots(s: str) -> int:
        n = 0
        for ch in s:
            if ch == ".":
                n += 1
            else:
                break
        return n

    def _resolve_crossref_target(target: str, current_path: str) -> str:
        if not target.startswith("."):
            return target

        dot_count = _count_leading_dots(target)
        remainder = target[dot_count:]
        current_parts = current_path.split(".")

        up_levels = max(dot_count - 1, 0)
        base_parts = current_parts[: max(0, len(current_parts) - up_levels)]

        if remainder:
            return ".".join([*base_parts, *remainder.split(".")])

        return ".".join(base_parts)

    def _find_module_prefix(identifier: str) -> str:
        parts = identifier.split(".")

        for i in range(len(parts), 0, -1):
            candidate = ".".join(parts[:i])
            try:
                if importlib.util.find_spec(candidate) is not None:
                    return candidate
            except (ImportError, ValueError, AttributeError, ModuleNotFoundError):
                continue

        if len(parts) > 1:
            return ".".join(parts[:-1])
        return identifier

    def _identifier_to_url(identifier: str) -> str:
        module_prefix = _find_module_prefix(identifier)
        page_path = module_prefix.replace(".", "/")
        return f"{docs_base_url}/{page_path}/#{identifier}"

    def _rewrite_crossrefs_in_text(text: str, current_path: str) -> str:
        def replacer(match: re.Match[str]) -> str:
            label = match.group(1)
            raw_target = match.group(2)
            resolved_target = _resolve_crossref_target(raw_target, current_path)
            url = _identifier_to_url(resolved_target)
            return f"[{label}]({url})"

        return _CROSSREF_RE.sub(replacer, text)

    def _rewrite_schema_node(
        node: Any,
        *,
        current_path: str,
        def_registry: dict[str, str],
    ) -> None:
        if isinstance(node, dict):
            description = node.get("description")
            if isinstance(description, str):
                node["description"] = _rewrite_crossrefs_in_text(
                    description,
                    current_path,
                )

            defs = node.get("$defs")
            if isinstance(defs, dict):
                for def_name, def_schema in defs.items():
                    nested_path = def_registry.get(def_name, current_path)
                    _rewrite_schema_node(
                        def_schema,
                        current_path=nested_path,
                        def_registry=def_registry,
                    )

            for key, value in node.items():
                if key == "$defs":
                    continue
                _rewrite_schema_node(
                    value,
                    current_path=current_path,
                    def_registry=def_registry,
                )

        elif isinstance(node, list):
            for item in node:
                _rewrite_schema_node(
                    item,
                    current_path=current_path,
                    def_registry=def_registry,
                )

    def rewrite_json_schema_crossrefs(
        schema: dict[str, Any],
        model: type[BaseModel] | None = None,
    ) -> None:
        if model is None:
            return

        root_path = f"{model.__module__}.{model.__name__}"
        def_registry = _build_def_registry(model)

        _rewrite_schema_node(
            schema,
            current_path=root_path,
            def_registry=def_registry,
        )

    return rewrite_json_schema_crossrefs
