"""A `griffe` extension to dynamically collect `isqx` definitions/details.

[Details][isqx.details.Details] are injected into the docstrings of attributes
and functions in the documentation generated by `mkdocstrings-python`.

!!! note

    You must install `isqx` with the `docs` extra optional dependencies
    to use this module.
"""

from __future__ import annotations

import ast
import inspect
import json
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, NamedTuple, Union

from griffe import (
    Alias,
    Docstring,
    Extension,
    Inspector,
    Object,
    ObjectNode,
    Visitor,
    dynamic_import,
    get_logger,
)
from griffe._internal.expressions import (
    Expr,
    ExprAttribute,
    ExprCall,
    ExprDict,
    ExprKeyword,
    ExprName,
    ExprSubscript,
    ExprTuple,
)
from griffe._internal.models import Attribute
from typing_extensions import TypeAlias

from .. import (
    Anchor,
    BasicFormatter,
    Dimensionless,
    LazyProduct,
    QtyKind,
    StrFragment,
    Tagged,
    kind,
    module_attribute_metadata,
)
from .. import (
    Expr as IsqxExpr,
)
from .._core import _ARGS_EXPR, _ARGS_NUMBER
from ..details import (
    _ARGS_DETAIL_KEY,
    Equation,
    HasKaTeXWhere,
    Symbol,
    Wikidata,
    _RefSelf,
)
from . import PATH_PLUGIN

if TYPE_CHECKING:
    from typing import Generator

    from griffe import GriffeLoader, Module
    from griffe._internal.models import Attribute

    from .. import AnnotatedMetadata, NamedExpr, Number
    from .._fmt import _BasicFormatterState
    from ..details import Detail, DetailKey, Details, RefDetail, WhereValue
    from .plugin import IsqxPluginConfig


logger = get_logger(__name__)


def get_templates_path() -> Path:
    """See: https://mkdocstrings.github.io/usage/handlers/#handler-extensions"""
    return PATH_PLUGIN / "templates"


_ARGS_DEFINITION = (LazyProduct, *_ARGS_NUMBER, *_ARGS_EXPR, QtyKind)
# NOTE: equilibrium constants are factory functions that return a QtyKind.
# right now we skip over them!

Unit: TypeAlias = Union[tuple[StrFragment, ...], None]
"""The unit of measurement *rendered* as string fragments."""
# Mul(M, S) is represented as (Anchor("M", "isqx.M"), "\cdot", Anchor("S", "isqx.S"))
# in the future, we want to return a tagged union (the direct representation
# of the unit). for now, this is good enough.
# we also want to get rid of the None.


# see: https://mkdocstrings.github.io/griffe/guide/users/how-to/selectively-inspect/
@dataclass(frozen=True)
class Where:
    symbol: str
    description: StrFragment | tuple[StrFragment, ...]
    unit: Unit  # for now


@dataclass(frozen=True)
class KaTeXWhere:
    katex: str
    where: tuple[Where, ...] | None = None


@dataclass(frozen=True)
class SymbolDetail(KaTeXWhere):
    remarks: str | None = None


@dataclass(frozen=True)
class EquationDetail(KaTeXWhere):
    assumptions: tuple[StrFragment | tuple[StrFragment, ...], ...] | None = None


@dataclass(frozen=True)
class WikidataDetail:
    qcode: str


@dataclass(frozen=True)
class Quantity:
    value: str  # for now: in the future we should support lazy products
    unit: Unit  # for now


CanonicalPath: TypeAlias = str


@dataclass
class QtyKindDetail:
    """Several quantities can share the same unit. Two cases:

    - Explicit inheritance: when a quantity is defined by subscripting
        another, like `POTENTIAL_ENERGY = ENERGY["potential"]`: this is a
        strong signal of a parent-child relationship.
        we can detect this during the griffe AST walking
    - Implicit grouping: some quantities are related by their physical
        dimension but not defined via subscripting (e.g., SPEED and VELOCITY).
        we can group them by their underlying SI dimension.
        the hierarchy can be inferred from the tags.
    """

    parent: str | None = None
    unit_si_coherent: Unit = None  # str for now
    tags: tuple[str, ...] = field(default_factory=tuple)  # str for now
    wikidata: list[WikidataDetail] = field(default_factory=list)
    symbols: list[SymbolDetail] = field(default_factory=list)
    equations: list[EquationDetail] = field(default_factory=list)


class IsqxExtension(Extension):
    def __init__(
        self, config: IsqxPluginConfig, objects_out_path: str | None = None
    ):
        self.config = config
        self.definitions: Definitions = {}
        self.objects: dict[str, QtyKindDetail] = {}  # by path
        self.objects_out_path = objects_out_path
        self.possible_parent_maps: dict[str, str] = {}  # child -> parent
        """Note that items in this map are not guaranteed to be isqx objects
        because griffe works with static analysis."""

    def on_instance(
        self,
        *,
        node: ast.AST | ObjectNode,
        obj: Object,
        agent: Visitor | Inspector,
        **kwargs: Any,
    ) -> None:
        inject_citation_into_docstring(obj, agent)

        # collect strict parent-child relationships
        if (
            isinstance(obj, Attribute)
            and isinstance(obj.value, ExprSubscript)
            and isinstance(
                parent_expr := obj.value.left, (ExprName, ExprAttribute)
            )
        ):
            parent_path = parent_expr.canonical_path
            if parent_path.split(".", 1)[0] in ("typing", "typing_extensions"):
                # heuristic: we captured something invalid like Union[str, int]
                return
            self.possible_parent_maps[obj.canonical_path] = parent_path

    def on_module_instance(
        self,
        *,
        node: ast.AST | ObjectNode,
        mod: Module,
        agent: Visitor | Inspector,
        **kwargs: Any,
    ) -> None:
        """Populate definitions from the current module being visited using
        runtime analysis, if it is matches the path under the configured
        [definitions][isqx.mkdocs.plugin.IsqxPluginConfig.definitions].

        Note that if module `a` contains `FOO: Annotated[float, isqx.M] = 1.0`
        and module `b` contains `from a import FOO`, then `self.definitions`
        will contain both (which points to the same object):

        - `a.FOO`: Definition(value=1.0, annotated_metadata=M)
        - `b.FOO`: Definition(value=1.0, annotated_metadata=None)

        which is not ideal. When `griffe` visits `b`, `on_alias` will make sure
        the definition for `b.FOO` is removed.
        """
        if mod.path not in self.config.definitions:
            return
        module_rt = dynamic_import(mod.path)
        metadata = dict(module_attribute_metadata(module_rt))
        self.definitions |= {
            f"{mod.path}.{name}": Definition(obj, metadata.get(name))
            for name, obj in inspect.getmembers(module_rt)
            if isinstance(obj, _ARGS_DEFINITION)
        }

    def on_alias_instance(
        self,
        *,
        node: ast.AST | ObjectNode,
        alias: Alias,
        agent: Visitor | Inspector,
        **kwargs: Any,
    ) -> None:
        # todo: actually we probably shouldn't do this
        if alias.path in self.definitions:
            del self.definitions[alias.path]

    def on_package(
        self, *, pkg: Module, loader: GriffeLoader, **kwargs: Any
    ) -> None:
        logger.info(f"loaded {len(self.definitions)} definitions")
        reverse_definitions = _build_reverse_definitions(self.definitions)
        formatter = MkdocsFormatter(reverse_definitions)

        for defs_path in self.config.details:
            for path, item in _process_details_module(
                defs_path, loader, formatter, reverse_definitions
            ):
                objects = self.objects.setdefault(path, QtyKindDetail())
                attr_target: Attribute = loader.modules_collection[path]
                isqx_extras = _get_or_init_isqx_extras(attr_target)

                if isinstance(item, WikidataDetail):
                    objects.wikidata.append(item)
                    isqx_extras["wikidata"].append(item)
                elif isinstance(item, SymbolDetail):
                    objects.symbols.append(item)
                    isqx_extras["symbols"].append(item)
                elif isinstance(item, EquationDetail):
                    objects.equations.append(item)
                    isqx_extras["details"].append(item)

        # set the parent (if found). would be nice for jinja to also show this
        # but we focus on the frontend vis for now.
        for child_path, parent_path in self.possible_parent_maps.items():
            if child_path in self.objects and parent_path in self.objects:
                child_object_detail = self.objects[child_path]
                child_object_detail.parent = parent_path

        for path, details in self.objects.items():
            if path not in self.definitions:
                continue
            definition = self.definitions[path]
            if not isinstance(v := definition.value, QtyKind):
                continue
            details.unit_si_coherent = tuple(
                formatter.fmt(v.unit_si_coherent)
            )  # str for now
            if v.tags is not None:
                details.tags = tuple(str(t) for t in v.tags)  # str for now

        if self.objects_out_path:
            constants = {}
            for path, definition in self.definitions.items():
                if not isinstance(
                    definition.value, (LazyProduct, *_ARGS_NUMBER)
                ):
                    continue
                value = definition.value
                value_str = str(
                    value.to_approx()
                    if isinstance(value, LazyProduct)
                    else value
                )
                unit_expr: IsqxExpr | None = None
                if not (
                    (meta := definition.annotated_metadata)
                    and (unit_expr := meta.unit)
                ):
                    logger.warning(
                        f"no unit found for {path} in annotated metadata"
                    )
                    continue  # dimensionless quantities should be annotated anyways
                constants[path] = Quantity(
                    value=value_str,
                    unit=(
                        tuple(formatter.fmt(unit_expr))  # str for now
                        if unit_expr is not None
                        else None
                    ),
                )

            json_path = Path(self.objects_out_path) / "objects.json"
            # do not serialize members that have `None` or empty tuple values to
            # save space
            output_data = {
                "qtyKinds": {
                    path: to_dict(details)
                    for path, details in self.objects.items()
                },
                "constants": {
                    path: to_dict(details)
                    for path, details in constants.items()
                },
                "units": {},  # TODO
            }
            json_path.parent.mkdir(parents=True, exist_ok=True)
            with open(json_path, "w") as f:
                json.dump(output_data, f)
            logger.info(f"wrote objects to {json_path}")


def to_dict(obj: Any) -> dict[str, Any]:
    return asdict(obj, dict_factory=lambda x: {k: v for (k, v) in x if v})


class Definition(NamedTuple):
    value: LazyProduct | Number | IsqxExpr | QtyKind
    annotated_metadata: AnnotatedMetadata | None


Definitions: TypeAlias = dict[CanonicalPath, Definition]


class _ReverseDefinition(NamedTuple):
    path: CanonicalPath
    annotated_metadata: AnnotatedMetadata | None


_ReverseDefinitions: TypeAlias = dict[int, _ReverseDefinition]
"""A mapping of the `builtins.id(value)` to the reverse definition."""
#
# unit formatting
#


class MkdocsFormatter(BasicFormatter):
    def __init__(
        self,
        reverse_definitions: _ReverseDefinitions,
    ):
        self.reverse_definitions = reverse_definitions
        super().__init__(verbose=False)

    def visit_named(
        self, expr: NamedExpr, state: _BasicFormatterState
    ) -> Generator[StrFragment, None, None]:
        name_formatted: str = next(super().visit_named(expr, state))  # type: ignore
        if reverse_def := self.reverse_definitions.get(id(expr)):
            yield Anchor(
                text=name_formatted,
                path=reverse_def.path,
            )
        else:
            yield name_formatted


#
# details
#


def parse_katex_where_static(
    katex_where_rt: HasKaTeXWhere,
    katex_where_st: str | Expr,
) -> tuple[str, ExprDict | None]:
    # working with static analysis is very annoying, we have to handle kwargs
    # and defaults. using a bunch of asserts for now.
    assert isinstance(katex_where_st, ExprCall), katex_where_st
    katex_st = katex_where_st.arguments[0]
    katex_st_value = (
        katex_st.value if isinstance(katex_st, ExprKeyword) else katex_st
    )
    assert isinstance(katex_st_value, str), katex_st_value
    katex_st_value = (
        katex_st_value.removeprefix("'")
        .removeprefix('"')
        .removesuffix("'")
        .removesuffix('"')
        .replace("\\\\", "\\")
    )  # regex is probably better but eh
    if not len(katex_where_st.arguments) > 1:
        return katex_st_value, None
    where_st = katex_where_st.arguments[1]
    where_st_value = where_st
    if isinstance(where_st, ExprKeyword):
        if where_st.name != "where":
            return katex_st_value, None
        where_st_value = where_st.value
    assert isinstance(where_st_value, ExprDict), where_st_value
    assert katex_where_rt.katex == katex_st_value, (
        f"`griffe` skipped an entry!"
        "\nthis can happen in two keys in the dictionary refer to the same "
        "value. check the definitions."
        f"\nruntime value {katex_where_rt.katex}\ndoes not match griffe: {katex_st_value}"
    )
    return katex_st_value, where_st_value


def parse_where(
    key_rt: DetailKey | Callable[..., DetailKey],
    detail_rt: HasKaTeXWhere,
    where_st: ExprDict,
    *,
    self_path: str,
    reverse_definitions: _ReverseDefinitions,
    formatter: MkdocsFormatter,
    mut_self_symbols: list[str],
) -> Generator[Where, None, None]:
    if detail_rt.where is None:
        return
    for (w_rt_k, w_rt_v), w_st_v in zip(
        detail_rt.where.items(),
        where_st.values,
    ):
        description = _parse_fragments(w_rt_v, w_st_v, self_path)
        unit_expr: IsqxExpr | None = None
        frags_for_unit_inference = (
            w_rt_v if isinstance(w_rt_v, tuple) else (w_rt_v,)
        )
        for frag in frags_for_unit_inference:
            if isinstance(frag, (str, Anchor)):
                continue
            if isinstance(frag, _RefSelf):
                mut_self_symbols.append(w_rt_k)
                if isinstance(key_rt, _ARGS_DETAIL_KEY):
                    unit_expr = _get_unit_expr(key_rt, reverse_definitions)
            else:
                unit_expr = _get_unit_expr(frag, reverse_definitions)
            if unit_expr is not None:
                break  # only infer unit from the first fragment

        unit: Unit
        if unit_expr is None:
            unit = None
        elif kind(unit_expr) == "dimensionless":
            unit = ("dimensionless",)
        else:
            unit = tuple(s for s in formatter.fmt(unit_expr))  # str  for now

        yield Where(
            symbol=w_rt_k,
            description=description,
            unit=unit,
        )


def _build_reverse_definitions(definitions: Definitions) -> _ReverseDefinitions:
    reverse_definitions: _ReverseDefinitions = {}
    for path, definition in definitions.items():
        i = id(definition.value)
        if i in reverse_definitions:
            logger.warning(
                f"duplicate definition id for {path} and {reverse_definitions[i].path}"
            )
            continue
        reverse_definitions[i] = _ReverseDefinition(
            path, definition.annotated_metadata
        )
    return reverse_definitions


def _get_or_init_isqx_extras(attr_target: Attribute) -> dict[str, Any]:
    mkdocstrings = attr_target.extra.setdefault("mkdocstrings", {})
    tmpl = "function" if attr_target.is_function else "attribute"
    mkdocstrings.setdefault("template", f"isqx/{tmpl}.html.jinja")
    isqx: dict[str, Any] = mkdocstrings.setdefault("isqx", {})
    isqx.setdefault("details", [])
    isqx.setdefault("symbols", [])
    isqx.setdefault("wikidata", [])
    return isqx


def _process_katex_where_detail(
    detail_rt: HasKaTeXWhere,
    detail_st: str | Expr,
    key_rt: DetailKey | Callable[..., DetailKey],
    self_path: str,
    formatter: MkdocsFormatter,
    reverse_definitions: _ReverseDefinitions,
) -> tuple[SymbolDetail | EquationDetail, list[str]]:
    _katex_st_value, where_st_value = parse_katex_where_static(
        detail_rt, detail_st
    )
    self_symbols: list[str] = []
    where = (
        tuple(
            parse_where(
                key_rt,
                detail_rt,
                where_st_value,
                self_path=self_path,
                formatter=formatter,
                reverse_definitions=reverse_definitions,
                mut_self_symbols=self_symbols,
            )
        )
        if where_st_value is not None
        else None
    )
    if isinstance(detail_rt, Equation):
        return EquationDetail(
            katex=detail_rt.katex,
            where=where,
            assumptions=(
                tuple(a) if (a := detail_rt.assumptions) is not None else None
            ),
        ), self_symbols
    elif isinstance(detail_rt, Symbol):
        return SymbolDetail(
            katex=detail_rt.katex,
            where=where,
            remarks=detail_rt.remarks,
        ), self_symbols
    raise ValueError(f"unknown {type(detail_rt)=}")


def _process_definitions(
    defs_rt: Detail | tuple[Detail, ...],
    defs_st: Expr | str,
    key_rt: DetailKey | Callable[..., DetailKey],
    self_path: str,
    formatter: MkdocsFormatter,
    reverse_definitions: _ReverseDefinitions,
) -> Generator[WikidataDetail | SymbolDetail | EquationDetail, None, None]:
    definitions_rt = defs_rt if isinstance(defs_rt, tuple) else (defs_rt,)
    definitions_st = (
        defs_st.elements if isinstance(defs_st, ExprTuple) else (defs_st,)
    )

    for detail_rt, detail_st in zip(definitions_rt, definitions_st):
        if isinstance(detail_rt, Wikidata):
            yield WikidataDetail(qcode=detail_rt.qcode)
        elif isinstance(detail_rt, HasKaTeXWhere):
            item, self_symbols = _process_katex_where_detail(
                detail_rt,
                detail_st,
                key_rt,
                self_path,
                formatter,
                reverse_definitions,
            )
            yield item
            if self_symbols:
                if len(self_symbols) > 1:
                    raise ValueError(
                        f"too many self symbols for {detail_rt.katex=}"
                    )
                yield SymbolDetail(katex=self_symbols[0])
        else:
            raise ValueError(f"unknown {type(detail_rt)=}")


def _process_details_module(
    defs_path: str,
    loader: GriffeLoader,
    formatter: MkdocsFormatter,
    reverse_definitions: _ReverseDefinitions,
) -> Generator[
    tuple[str, WikidataDetail | SymbolDetail | EquationDetail], None, None
]:
    details_rt: Details = dynamic_import(defs_path)
    details_st: Attribute = loader.modules_collection[defs_path]
    assert isinstance(details_st.value, ExprDict), details_st

    for (key_st, defs_st), (key_rt, defs_rt) in zip(
        zip(details_st.value.keys, details_st.value.values),
        details_rt.items(),
    ):
        assert isinstance(key_st, (ExprAttribute, ExprName)), key_st
        self_path = key_st.canonical_path
        for item in _process_definitions(
            defs_rt,
            defs_st,
            key_rt,
            self_path,
            formatter,
            reverse_definitions,
        ):
            yield (self_path, item)


def _parse_fragments(
    where_rt: WhereValue, where_st: Expr | str, self_path: str
) -> StrFragment | tuple[StrFragment, ...]:
    if isinstance(where_rt, tuple):
        fragments_rt = where_rt
        fragments_st = where_st.elements  # type: ignore
        parsed_frags = tuple(
            _parse_fragment(rt, st, self_path, i)
            for i, (rt, st) in enumerate(zip(fragments_rt, fragments_st))
        )
        return parsed_frags
    return _parse_fragment(where_rt, where_st, self_path, 0)


def _parse_fragment(
    frag_rt: WhereValue, frag_st: Expr | str, self_path: str, i: int
) -> StrFragment:
    if isinstance(frag_rt, (str, Anchor)):
        return frag_rt
    if isinstance(frag_rt, _RefSelf):
        return Anchor(screaming_to_normal(self_path, i), self_path)
    # RefDetail
    assert isinstance(frag_st, (ExprAttribute, ExprName)), (
        "expected value of the where clause to be a reference to a pre-defined "
        f"quantity, but got {frag_st} (an inline expression)\n"
        "= help: define the quantity first and then use it in the where clause"
    )
    path = frag_st.canonical_path
    return Anchor(screaming_to_normal(path, i), path)


def _get_unit_expr(
    meaning: RefDetail, reverse_definitions: _ReverseDefinitions
) -> IsqxExpr | None:
    if isinstance(meaning, QtyKind):
        return meaning.unit_si_coherent
    elif isinstance(meaning, (Dimensionless, Tagged)):
        return meaning
    elif isinstance(meaning, (LazyProduct, *_ARGS_NUMBER)):
        if (
            reverse_def := reverse_definitions.get(id(meaning))
        ) and reverse_def.annotated_metadata:
            return reverse_def.annotated_metadata.unit
        else:
            logger.warning(f"no unit found for {meaning} in registry")
            return None
    return None


def screaming_to_normal(s: str, i: int) -> str:
    # a better way would be to parse the docstring but that might be too verbose
    normalized = s.split(".")[-1].lower().replace("_", " ")
    return normalized.capitalize() if i == 0 else normalized


def inject_citation_into_docstring(
    obj: Object, agent: Visitor | Inspector
) -> None:
    if not obj.canonical_path.startswith("isqx._citations."):
        return
    # quick hack to put runtime value into the docstring
    doc_str = dynamic_import(obj.path)
    if not isinstance(doc_str, str):
        logger.warning(
            f"expected citation {obj.path} to be a string, got {type(doc_str)}"
        )
        return
    obj.docstring = Docstring(
        doc_str,
        parent=obj,
        parser=agent.docstring_parser,
        parser_options=agent.docstring_options,
    )
