import ast
import enum
import inspect
import typing
from datetime import date, datetime, time, timedelta
from decimal import Decimal
from types import NoneType
from typing import Any, Callable, Generator, Literal

import pydantic
from pydantic_core import PydanticUndefined
from typing_extensions import Self

from stable_pydantic.imports import Imports
from stable_pydantic.model_graph import ModelNode
from stable_pydantic.utils import BASE_TYPES, COMPOSITE_TYPES, UNION_TYPES, Unreachable


def clean(model: ModelNode) -> tuple[str, Imports, set[str]]:
    """Returns the cleaned source code, imports, and set of forward references."""

    tree = model.get_ast()
    if len(tree.body) != 1:
        raise ValueError("Model must have exactly one class definition.")
    if not isinstance(tree.body[0], ast.ClassDef):
        raise ValueError("Model must have a class definition.")
    if tree.body[0].decorator_list:
        raise ValueError("Decorators are not supported on models.")

    keep_only_class_fields(tree.body[0])

    current_field_count = len(model.node.__annotations__)
    ast_field_count = len(
        [v for v in tree.body[0].body if isinstance(v, ast.AnnAssign)]
    )
    if ast_field_count != current_field_count:
        raise ValueError(
            f"Field count mismatch. Expected {current_field_count}, got {ast_field_count}."
        )

    # Note that you may have inherited some fields from the parent, so this doesn't hold
    # assert current_field_count == len(model.node.model_fields)

    touched_base_types, forward_refs = _clean(model, tree.body[0])

    # Update the AST class name to the desired name
    tree.body[0].name = model.aliased_name()

    # Update base class names to use aliased names
    _clean_bases(model, tree.body[0])

    source = ast.unparse(tree)
    return source, touched_base_types, forward_refs


def _clean_bases(model: ModelNode, class_def: ast.ClassDef):
    """Update base class references to use aliased names."""
    if not model.bases():
        return

    for base in class_def.bases:
        assert isinstance(base, (ast.Name, ast.Attribute))

    # Replace all base classes with their aliased names
    class_def.bases = [
        ast.Name(id=base.aliased_name(), ctx=ast.Load()) for base in model.bases()
    ]


def keep_only_class_fields(class_def: ast.ClassDef):
    new_body = []
    for node in class_def.body:
        if isinstance(node, ast.AnnAssign):
            # Field with type annotation like `name: str` or `name: str = value`
            new_body.append(node)
        elif (
            isinstance(node, ast.Assign)
            and len(node.targets) == 1
            and isinstance(node.targets[0], ast.Name)
            and node.targets[0].id == "model_config"
        ):
            # Handle model_config = ConfigDict(...)
            # (No annotation)
            new_body.append(node)
        elif isinstance(node, ast.Pass):
            continue
        elif isinstance(node, ast.FunctionDef):
            continue
        elif (
            isinstance(node, ast.Expr)
            and isinstance(node.value, ast.Constant)
            and isinstance(node.value.value, str)
        ):
            # Skip docstrings
            continue
        else:
            raise ValueError(f"Unsupported node: {node}\n{ast.dump(node)}")

    if not new_body:
        # If no fields, add pass statement
        new_body: list[ast.stmt] = [ast.Pass()]

    class_def.body = new_body


def _clean(model: ModelNode, tree: ast.ClassDef) -> tuple[Imports, set[str]]:
    imports = Imports.empty()
    forward_refs = set()

    hints = typing.get_type_hints(model.node)  # Resolves forward references

    for raw in tree.body:
        if isinstance(raw, ast.Pass):
            continue

        if (
            isinstance(raw, ast.Assign)
            and len(raw.targets) == 1
            and isinstance(raw.targets[0], ast.Name)
            and raw.targets[0].id == "model_config"
        ):
            # model_config = ConfigDict(...)
            # (No annotation)
            imports.config_dict = True
            _clean_model_config(raw)
            continue

        assert isinstance(raw, ast.AnnAssign)
        ast_node: ast.AnnAssign = raw
        assert isinstance(ast_node.target, ast.Name)

        field_name: str = ast_node.target.id
        if field_name == "model_config":
            imports.config_dict = True
            _clean_model_config(ast_node)
            continue

        model_field = model.node.model_fields[field_name]

        def replace(a: ast.expr):
            ast_node.annotation = a

        _clean_field_type(
            hints.get(field_name),
            ast_node.annotation,
            replace,
            lambda: forward_refs.add(field_name),
            model.all,
            imports,
        )

        if isinstance(ast_node.value, ast.Call):
            call: ast.Call = ast_node.value

            direct = isinstance(call.func, ast.Name) and call.func.id == "Field"
            indirect = (
                isinstance(call.func, ast.Attribute) and call.func.attr == "Field"
            )

            if direct or indirect:
                imports.field = True

            if indirect:
                ast_node.value = ast.Call(
                    func=ast.Name(id="Field", ctx=ast.Load()),
                    args=call.args,
                    keywords=call.keywords,
                )

        if ast_node.value is not None:
            cleaned_default = _clean_field_default(
                model_field.get_default(call_default_factory=True),
                ast_node.value,
                model.all,
                imports,
            )
            if cleaned_default is not None:
                ast_node.value = cleaned_default

    return imports, forward_refs


def _clean_field_type(
    field_type: type[Any] | Any | None,
    annotation: ast.expr,
    replace: Callable[[ast.expr], None],
    # Could be a return type, but then it's easy to miss.
    is_forward_ref: Callable[[], None],
    all: dict[type[pydantic.BaseModel], ModelNode],
    imports: Imports,
):
    if field_type is None:
        raise ValueError(f"Field type is None: {annotation}")

    elif field_type in BASE_TYPES:
        imports.touch(field_type)
        # TODO: assert the types match.
        pass

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

    elif isinstance(annotation, ast.Constant) and isinstance(annotation.value, str):
        parsed = ast.parse(annotation.value, mode="eval")

        def replace_inner(a: ast.expr):
            parsed.body = a

        _clean_field_type(
            field_type, parsed.body, replace_inner, is_forward_ref, all, imports
        )
        annotation.value = ast.unparse(parsed.body)
        is_forward_ref()

    elif typing.get_origin(field_type) in UNION_TYPES:
        field_types = list(typing.get_args(field_type))
        flattened_annotation = _flatten_union_annotation(
            annotation, imports, has_none=NoneType in field_types
        )

        if len(field_types) != len(flattened_annotation):
            raise ValueError(
                "Union type length mismatch. This happens when different ways of writing union types are mixed.\n"
                + "Please use a single one at a time of `int | None` or `Union[int, None]` or `Optional[int]`."
            )

        for inner_field_type, (inner_annotation, replace) in zip(
            field_types, flattened_annotation, strict=True
        ):
            _clean_field_type(
                inner_field_type,
                inner_annotation,
                replace,
                is_forward_ref,
                all,
                imports,
            )

    elif typing.get_origin(field_type) is Literal:
        imports.literal = True

        values = list(typing.get_args(field_type))

        assert isinstance(annotation, ast.Subscript)
        annotations = _unpack_composite_annotation(annotation)

        for inner_value, (inner_annotation, replace) in zip(
            values, annotations, strict=True
        ):
            assert type(inner_value) in BASE_TYPES
            assert isinstance(inner_annotation, ast.Constant)
            assert inner_annotation.value == inner_value

    elif typing.get_origin(field_type) in COMPOSITE_TYPES:
        imports.touch(typing.get_origin(field_type))  # type: ignore[arg-type]

        field_types = list(typing.get_args(field_type))

        assert isinstance(annotation, ast.Subscript)
        annotations = _unpack_composite_annotation(annotation)

        for inner_field_type, (inner_annotation, replace) in zip(
            field_types, annotations, strict=True
        ):
            _clean_field_type(
                inner_field_type,
                inner_annotation,
                replace,
                is_forward_ref,
                all,
                imports,
            )

    elif field_type is Self:
        imports.self_type = True
        pass

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

    elif inspect.isclass(field_type) and issubclass(field_type, pydantic.BaseModel):
        inline_model: ModelNode = all[field_type]
        desired_name = inline_model.aliased_name()

        if isinstance(annotation, ast.Name):
            # Update annotation so it matches the expected name.
            annotation.id = desired_name
        elif isinstance(annotation, ast.Attribute):
            # Or overwrite it entirely if it's an attribute (field: some.attr.Path)
            replace(ast.Name(id=desired_name, ctx=annotation.ctx))
        else:
            raise ValueError(
                f"Expected name or attribute annotation, got: {ast.dump(annotation)}"
            )

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

    return


def _clean_field_default(
    default: Any,
    value: ast.expr,
    all: dict[type[pydantic.BaseModel], ModelNode],
    imports: Imports,
) -> ast.expr | None:
    """Clean field default values by serializing to JSON and replacing with json.loads()."""
    if default is not PydanticUndefined:
        actual_default = default
    else:
        return None

    if isinstance(
        actual_default, (type(None), bool, int, float, str, bytes, complex, Decimal)
    ):
        new_value = ast.Constant(value=actual_default)  # type: ignore[arg-type]
    elif isinstance(actual_default, (datetime, date, time, timedelta)):
        stringified_constant = ast.unparse(ast.Constant(value=actual_default))  # type: ignore[arg-type]
        assert stringified_constant.startswith("datetime.")
        reparsed_module = ast.parse(stringified_constant[len("datetime.") :])
        assert len(reparsed_module.body) == 1
        reparsed_constant = reparsed_module.body[0].value  # type: ignore
        assert isinstance(reparsed_constant, ast.Call)
        new_value = reparsed_constant
    elif isinstance(actual_default, pydantic.BaseModel):
        # For Pydantic models, serialize to JSON and load inline with model_validate_json
        model_type = type(actual_default)
        model_node = all.get(model_type)
        assert model_node is not None, (
            f"Model type {model_type} not found in the model tree."
        )

        # ModelName.model_validate_json("...")
        new_value = ast.Call(
            func=ast.Attribute(
                value=ast.Name(id=model_node.aliased_name(), ctx=ast.Load()),
                attr="model_validate_json",
                ctx=ast.Load(),
            ),
            args=[ast.Constant(value=actual_default.model_dump_json())],
            keywords=[],
        )
    else:
        imports.uses_json = True

        # We go through pydantic because it supports more types than json.dumps.
        default_str: str = (
            pydantic.TypeAdapter(type(actual_default))
            .dump_json(actual_default)
            .decode()
        )

        # json.loads("...")
        new_value = ast.Call(
            func=ast.Attribute(
                value=ast.Name(id="json", ctx=ast.Load()),
                attr="loads",
                ctx=ast.Load(),
            ),
            args=[ast.Constant(value=default_str)],
            keywords=[],
        )

    # Check if value is a Field() call, and if so update its arguments
    if isinstance(value, ast.Call) and (
        (isinstance(value.func, ast.Attribute) and value.func.attr == "Field")
        or (isinstance(value.func, ast.Name) and value.func.id == "Field")
    ):
        # Replace 'default' or 'default_factory' keyword arguments
        modified = False

        for keyword in value.keywords:
            if keyword.arg == "default_factory":
                if modified:
                    raise ValueError("Multiple default or default_factory keywords")
                modified = True

                keyword.value = ast.Lambda(
                    args=ast.arguments(
                        posonlyargs=[],
                        args=[],
                        kwonlyargs=[],
                        kw_defaults=[],
                        defaults=[],
                    ),
                    body=new_value,
                )
            elif keyword.arg == "default":
                if modified:
                    raise ValueError("Multiple default or default_factory keywords")
                modified = True

                keyword.value = new_value

        if not modified:
            raise ValueError("No default or default_factory keyword found")

        return value
    else:
        return new_value


def _unpack_composite_annotation(
    annotation: ast.Subscript,
) -> list[tuple[ast.expr, Callable[[ast.expr], None]]]:
    """Extract inner type annotations from subscripted composite types like list[int] or dict[str, int]."""
    slice_node = annotation.slice
    if isinstance(slice_node, ast.Tuple):
        # Multiple type parameters like dict[str, int] or tuple[int, str]

        def replace_multi(a: ast.expr, i: int):
            slice_node.elts[i] = a

        return [
            (elt, lambda a, i=i: replace_multi(a, i))
            for i, elt in enumerate(slice_node.elts)
        ]
    else:
        # Single type parameter like list[int] or set[str]
        def replace(a: ast.expr):
            annotation.slice = a

        return [(slice_node, replace)]


def _flatten_union_annotation(
    annotation: ast.expr, imports: Imports, has_none: bool
) -> list[tuple[ast.expr, Callable[[ast.expr], None]]]:
    def raise_unreachable(_: ast.expr) -> None:
        raise Unreachable()

    if isinstance(annotation, ast.Subscript):
        unpacked = _unpack_composite_annotation(annotation)

        if isinstance(annotation.value, ast.Name) and annotation.value.id == "Optional":
            imports.optional = True

            assert len(unpacked) == 1
            assert has_none

            unpacked, replace = unpacked[0]
            none = (ast.Constant(value=None), raise_unreachable)
            return [(unpacked, replace), none]

        elif isinstance(annotation.value, ast.Name) and annotation.value.id == "Union":
            imports.union = True
            return unpacked

        else:
            raise ValueError(f"Unsupported union type {annotation}")

    elif isinstance(annotation, ast.BinOp):
        unpacked = list(_iter_binop(annotation))
        return unpacked

    else:
        raise ValueError(f"Unsupported union type {annotation}")


def _iter_binop(
    node: ast.BinOp,
) -> Generator[tuple[ast.expr, Callable[[ast.expr], None]], None, None]:
    def replace_left(a: ast.expr):
        node.left = a

    def replace_right(a: ast.expr):
        node.right = a

    if isinstance(node.left, ast.BinOp):
        assert isinstance(node.left.op, ast.BitOr)
        yield from _iter_binop(node.left)
    else:
        yield node.left, replace_left

    yield node.right, replace_right


def _clean_model_config(
    node: ast.AnnAssign | ast.Assign,
) -> None:
    """Replace calls to `pydantic.ConfigDict` with a call to `ConfigDict`."""

    config = node.value
    if config is None:
        raise ValueError("model_config is missing")

    if not isinstance(config, ast.Call):
        raise ValueError(
            f"Unsupported model_config: expected ConfigDict(...), got {ast.unparse(config)}"
        )

    if isinstance(config.func, ast.Name) and config.func.id == "ConfigDict":
        # Already `ConfigDict(...)`, nothing to do
        pass
    elif isinstance(config.func, ast.Attribute) and config.func.attr == "ConfigDict":
        # `pydantic.ConfigDict(...)` -> `ConfigDict(...)`
        config.func = ast.Name(id="ConfigDict", ctx=config.func.ctx)
    else:
        raise ValueError(
            f"Unsupported model_config: expected ConfigDict(...), got {ast.unparse(config)}"
        )

    # Also clean the annotation if present
    if isinstance(node, ast.AnnAssign):
        ann = node.annotation
        if isinstance(ann, ast.Name) and ann.id == "ConfigDict":
            # Already `ConfigDict`, nothing to do
            pass
        elif isinstance(ann, ast.Attribute) and ann.attr == "ConfigDict":
            # `pydantic.ConfigDict` -> `ConfigDict`
            node.annotation = ast.Name(id="ConfigDict", ctx=ann.ctx)
        else:
            raise ValueError(
                f"Unsupported model_config annotation: expected ConfigDict, got {ast.unparse(ann)}"
            )
