import ast
import enum
import inspect
import sys
import types
import typing
from collections import OrderedDict
from graphlib import TopologicalSorter
from typing import Any, Literal

import pydantic
from typing_extensions import Self

from stable_pydantic.imports import Imports
from stable_pydantic.utils import (
    BASE_TYPES,
    COMPOSITE_TYPES_EXT,
    flatten,
)


class ModelNode(pydantic.BaseModel):
    node: type[pydantic.BaseModel]
    base_nodes: list[type[pydantic.BaseModel]]
    """Non-BaseModel parent classes this inherits from."""
    field_nodes: OrderedDict[str, list[type[pydantic.BaseModel]]]
    all: dict[type[pydantic.BaseModel], Self]

    alias: int | None = None
    """To deduplicate models that have the same name."""

    module: types.ModuleType

    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)

    @staticmethod
    def new(
        model: type[pydantic.BaseModel], module: types.ModuleType | None = None
    ) -> "ModelNode":
        if module is None:
            module = sys.modules[model.__module__]
        assert module is not None
        return _new(model, module)

    def bases(self) -> list[Self]:
        return [self.all[base] for base in self.base_nodes]

    def fields(self) -> OrderedDict[str, list[Self]]:
        return OrderedDict(
            (name, [self.all[field] for field in fields])
            for name, fields in self.field_nodes.items()
        )

    def ordered_nodes(self) -> list[Self]:
        """Return nodes in order (root first)."""

        def reverse(list: list[Any]) -> list[Any]:
            return list[::-1]

        graph: dict[type[pydantic.BaseModel], list[type[pydantic.BaseModel]]] = {
            k: reverse(v._edges()) for k, v in self.all.items()
        }

        ts = TopologicalSorter(graph)
        order = list(ts.static_order())
        order.reverse()

        return [self.all[node] for node in order]

    def clean_source_recursive(self) -> str:
        """Generate source code for a Pydantic model with only its fields (no methods/properties)."""

        nodes = self.ordered_nodes()
        nodes.reverse()
        last = nodes[-1] if nodes else None

        source = ""
        imports = Imports.empty()
        for node in nodes:
            new, new_touched_base_types = node.clean_source()
            source += new
            imports = imports.combine(new_touched_base_types)
            if node != last:
                source += "\n\n\n"

        return "# fmt: off\n" + imports.imports() + "\n\n\n" + source + "\n"

    def clean_source(self) -> tuple[str, Imports]:
        from stable_pydantic import clean

        source, imports, _forward_refs = clean.clean(self)

        # For compatibility across Python versions.
        source = source.replace("lambda :", "lambda:")

        return source, imports

    def get_ast(self) -> ast.Module:
        return ast.parse(inspect.getsource(self.node))

    def get_class_ast(self) -> ast.ClassDef:
        tree = self.get_ast()
        assert len(tree.body) == 1
        assert isinstance(tree.body[0], ast.ClassDef)
        return tree.body[0]

    def aliased_name(self) -> str:
        # TODO: rename anything that might collide with primitive types.
        if self.alias is None:
            return self.node.__name__
        else:
            return f"{self.node.__name__}_{self.alias}"

    def _all_forward_ref_models(self) -> set[type[pydantic.BaseModel]]:
        from stable_pydantic import clean

        _source, _imports, forward_refs_fields = clean.clean(self)

        forward_models: list[list[type[pydantic.BaseModel]]] = list(
            self.field_nodes[ref] for ref in forward_refs_fields
        )

        return set(flatten(forward_models))

    def _edges(self) -> list[type[pydantic.BaseModel]]:
        """Returns the children of the model, excluding forward references."""

        forward_refs = self._all_forward_ref_models()
        children: list[type[pydantic.BaseModel]] = []
        seen: set[type[pydantic.BaseModel]] = set()

        for base in self.base_nodes:
            if base not in forward_refs and base not in seen:
                children.append(base)
                seen.add(base)

        for field in self.field_nodes.values():
            for child in field:
                if child not in forward_refs and child not in seen:
                    children.append(child)
                    seen.add(child)

        return children


def _new(root: type[pydantic.BaseModel], module: types.ModuleType) -> ModelNode:
    def from_field_type(field_type: type[Any] | None) -> list[ModelNode] | None:
        if field_type is None or field_type in BASE_TYPES:
            return None

        elif field_type is enum.Enum:
            raise ValueError(f"Enum type {field_type} is todo")

        elif typing.get_origin(field_type) in COMPOSITE_TYPES_EXT:
            values = []
            for value in typing.get_args(field_type):
                new = from_field_type(value)
                if new:
                    values.extend(new)
            return values if values else None

        elif typing.get_origin(field_type) is Literal:
            values = list(typing.get_args(field_type))
            for value in values:
                assert type(value) in BASE_TYPES
            return None

        elif field_type is Self:
            return None

        elif field_type is ...:
            # Ellipsis in tuple[T, ...]
            return None

        elif inspect.isclass(field_type) and issubclass(field_type, pydantic.BaseModel):
            return [_new_inner(field_type)]

        else:
            raise ValueError(f"Unsupported type {field_type!r}")

    def _populate_aliases(self: ModelNode):
        all = {}

        for model, node in self.all.items():
            if model.__name__ not in all:
                all[model.__name__] = []

            all[model.__name__].append(node)

        duplicates = ((name, nodes) for name, nodes in all.items() if len(nodes) > 1)

        for _name, nodes in duplicates:
            for i, node in enumerate(nodes):
                node.alias = i

    all = {}

    def _new_inner(model: type[pydantic.BaseModel]) -> ModelNode:
        if model in all:
            return all[model]

        all[model] = _Marker.empty(model, module)

        hints = typing.get_type_hints(model, globalns=module.__dict__)

        children: OrderedDict[str, list[type[pydantic.BaseModel]]] = OrderedDict()
        for field_name in model.model_fields:
            field = from_field_type(hints.get(field_name))
            if field is None:
                continue
            else:
                children[field_name] = [node.node for node in field]

        # Track parent models (non-BaseModel bases)
        bases: list[type[pydantic.BaseModel]] = []
        for base in model.__bases__:
            if base is not pydantic.BaseModel and issubclass(base, pydantic.BaseModel):
                bases.append(_new_inner(base).node)

        node = ModelNode(
            node=model,
            base_nodes=bases,
            field_nodes=children,
            all={},  # Populated later
            module=module,
        )

        all[model] = node

        return node

    self = _new_inner(root)

    # Populate all
    assert all != {}
    for node in all.values():
        node.all = all

    _populate_aliases(self)

    return self


class _Marker(ModelNode):
    @staticmethod
    def empty(node: type[pydantic.BaseModel], module: types.ModuleType) -> ModelNode:
        return ModelNode(
            node=node,
            base_nodes=[],
            field_nodes=OrderedDict(),
            all={},
            module=module,
        )
