from __future__ import annotations

import ast
from pathlib import Path
from typing import Literal, Self, cast

from pydantic import BaseModel, Field, model_validator

from ..tools.filesystem import Dir, FileSystem
from ..util import filter as lfilter, find, maybe_or_call, not_none
from .base import Module, ModuleDep

NameAndDoc = tuple[str, str | None]
CodeElements = tuple[list[NameAndDoc], list[NameAndDoc], list[str], str | None]


class PythonModule(Module):
    """A Python module."""

    src_lang: Literal['python'] = 'python'
    imports: list[AstImport | AstImportFrom]
    elements: CodeElements

    @classmethod
    def from_content(cls, relative_path: Path, content: str) -> Self:
        try:
            elements = _parse_code_elements(ast.parse(content))
        except Exception as e:
            print(f'Failed to parse code elements from {relative_path}: {e}')
            elements = ([], [], [], None)

        return cls(
            relative_path=relative_path,
            content=content,
            imports=parse_imports(content),
            elements=elements,
        )

    @model_validator(mode='after')
    def gen_name(self) -> Self:
        self.name = rel_path_to_python_module_name(self.relative_path)
        return self


def _parse_code_elements(tree: ast.Module) -> CodeElements:
    """Parse functions, classes, constants from a Python file

    Returns:
        functions (list[tuple[str, str | None]]): functions and their docstrings
        classes (list[tuple[str, str | None]]): classes and their docstrings
        constants (list[str]): constants
        docstring (str | None): file-level docstring
    """
    functions = []
    classes = []
    constants = []
    docstring = None

    # fmt: off
    def maybe_docstring(node):
        def maybe_head(x): return x[0] if len(x) > 0 else None
        def value_of_expr(x): return x.value if isinstance(x, ast.Expr) else None
        def value_of_const(x): return x.value if isinstance(x, ast.Constant) else None
        def str_val(x): return x if isinstance(x, str) else None
        return str_val(value_of_const(value_of_expr(maybe_head(node.body))))
    # fmt: on

    # Get module docstring
    docstring = maybe_docstring(tree)

    def name_and_docstring(node):
        return node.name, maybe_docstring(node)

    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
            node = cast(ast.FunctionDef | ast.AsyncFunctionDef, node)
            functions.append(name_and_docstring(node))
        elif isinstance(node, ast.ClassDef):
            classes.append(name_and_docstring(node))
        elif isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id.isupper():
                    constants.append(target.id)

    return functions, classes, constants, docstring


def rel_path_to_python_module_name(path: Path) -> str:
    """
    Example:
    - `basic.py` -> `basic`
    - `utils/helpers.py` -> `utils.helpers`
    - `advanced/geometry/shapes.py` -> `advanced.geometry.shapes`
    """
    parts = list(path.parts)
    parts = [p.removesuffix('.py') for p in parts]
    return '.'.join(parts)


def build_py_dep_graph(fs: FileSystem) -> ModuleDep[PythonModule]:
    # Parse all Python modules
    py_modules: list[PythonModule] = []
    for node in fs.nodes:
        rel_path = node.relative_path
        if isinstance(node, Dir) or rel_path.suffix != '.py':
            continue
        py_modules.append(PythonModule.from_content(rel_path, node.content))

    adjacency_list: dict[str, set[str]] = {m.path_key: set() for m in py_modules}

    for py_module in py_modules:
        for import_ast in py_module.imports:
            targets = resolve_import(import_ast, py_module, py_modules, fs)

            for t in targets:
                if t != py_module:
                    adjacency_list[py_module.path_key].add(t.path_key)

    return ModuleDep.from_adjacency_list(py_modules, adjacency_list)


class AstImport(BaseModel):
    """Serializable ast.Import"""

    lineno: int = Field(description='Line number where import appears')
    col_offset: int = Field(description='Column offset where import appears')
    end_lineno: int | None = Field(
        default=None, description='Line number where import ends'
    )
    end_col_offset: int | None = Field(
        default=None, description='Column offset where import ends'
    )
    aliases: list[str] = Field(description='Aliases used in import')

    @classmethod
    def from_ast(cls, node: ast.Import) -> Self:
        return cls(
            lineno=node.lineno,
            col_offset=node.col_offset,
            end_lineno=node.end_lineno,
            end_col_offset=node.end_col_offset,
            aliases=[alias.name for alias in node.names],
        )


class AstImportFrom(AstImport):
    """Serializable ast.ImportFrom"""

    module: str | None = Field(description='Module being imported from')
    level: int = Field(description='Level of import (0=absolute, 1+=relative)')

    @classmethod
    def from_ast(cls, node: ast.ImportFrom) -> Self:
        return cls(
            lineno=node.lineno,
            col_offset=node.col_offset,
            end_lineno=node.end_lineno,
            end_col_offset=node.end_col_offset,
            aliases=[alias.name for alias in node.names],
            module=node.module,
            level=node.level,
        )


def parse_imports(content: str) -> list[AstImport | AstImportFrom]:
    """Parse import statements from a Python file"""
    imports = []
    tree = ast.parse(content)

    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            imports.append(AstImport.from_ast(node))

        elif isinstance(node, ast.ImportFrom):
            imports.append(AstImportFrom.from_ast(node))
    return imports


def resolve_import(
    import_ast: AstImport | AstImportFrom,
    importing_module: PythonModule,
    all_modules: list[PythonModule],
    fs: FileSystem,
) -> list[PythonModule]:
    """
    Resolve an import statement to 0+ target modules.
    """

    # fmt: off
    def map_filter(f, xs): return lfilter(not_none, map(f, xs))
    def path_key_is(p): return lambda m: m.path_key == p
    def find_module(p): return find(path_key_is(p), all_modules)
    # fmt: on

    def for_import_from(ast: AstImportFrom):
        # level 0: from pkg2 import b, c
        # level >0: from .foo import bar
        # fmt: off
        base = (ast.module if ast.level == 0 else
                resolve_relative_import_base_module(ast, importing_module))
        # fmt: on

        def import_from_alias(alias):
            # alias strings are names for ImportFrom (e.g. "b", or "*")
            name = alias.split('.')[0]

            # conservative: if star import, fall back to the base module/package itself
            # if resolved is None, fallback to depend on the base module/package
            # resolved_path is None or like "pkg2/b.py" or "pkg2/__init__.py"
            # fmt: off
            def find_base(): return find_module_by_name(cast(str, base), all_modules)
            return (find_base() if name == '*' else
                    maybe_or_call(find_base, find_module,
                                  resolve_from_import(fs, cast(str, base), name)))
            # fmt: on

        return map_filter(import_from_alias, import_ast.aliases) if base else []

    def for_plain_import():
        def for_import_alias(alias):
            # your AstImport stores alias.name strings
            # Prefer exact dotted match first (pkg2.b -> pkg2/b.py if present)
            dotted = alias.split('.')
            cand = str(Path(*dotted).with_suffix('.py'))
            # fmt: off
            mod_of_dotted = (None if fs.get_node_by_path(cand) is None else
                             find_module(cand))
            # fmt: on
            # or fall back to existing heuristics (root module / filename / etc)
            return mod_of_dotted or find_module_by_name(dotted[0], all_modules)

        return map_filter(for_import_alias, import_ast.aliases)

    # fmt: off
    return (for_import_from(import_ast) if isinstance(import_ast, AstImportFrom) else
            for_plain_import() if isinstance(import_ast, AstImport) else [])
    # fmt: on


def resolve_relative_import_base_module(
    import_ast: AstImportFrom,
    importing_module: PythonModule,
) -> str | None:
    """
    Convert a relative import (level>0) into an absolute-ish dotted base module name.

    Examples (importing_module = pkg1/a.py => pkg1.a):
      from .foo import bar   -> base module "pkg1.foo"
      from ..pkg2 import b   -> base module "pkg2"
      from . import x        -> base module "pkg1"
    """
    # importing module name: pkg1.a
    importing_name = cast(str, importing_module.name)
    importing_pkg_parts = importing_name.split('.')[:-1]  # package parts only

    # level=1 means "current package"
    up = max(import_ast.level - 1, 0)
    if up > len(importing_pkg_parts):
        return None

    base_parts = importing_pkg_parts[: len(importing_pkg_parts) - up]

    if import_ast.module:
        base_parts += import_ast.module.split('.')

    if not base_parts:
        return None

    return '.'.join(base_parts)


def find_module_by_name(
    target_name: str,
    all_modules: list[PythonModule],
) -> PythonModule | None:
    """Find a module by name, trying different matching strategies."""

    # Try exact match first
    for module in all_modules:
        if module.name == target_name:
            return module

    # Try matching just the filename (stem)
    for module in all_modules:
        if module.relative_path.stem == target_name:
            return module

    # Try matching as root of dotted name
    for module in all_modules:
        module_name = cast(str, module.name)
        if module_name.split('.')[0] == target_name:
            return module

    return None


def resolve_from_import(fs, base_mod: str, name: str) -> str | None:
    # base_mod like "pkg2", name like "b"
    cand_mod = Path(*base_mod.split('.')) / f'{name}.py'
    if fs.get_node_by_path(str(cand_mod)) is not None:
        return str(cand_mod)

    cand_pkg = Path(*base_mod.split('.')) / name / '__init__.py'
    if fs.get_node_by_path(str(cand_pkg)) is not None:
        return str(cand_pkg)

    # fall back: package itself
    base_pkg_init = Path(*base_mod.split('.')) / '__init__.py'
    if fs.get_node_by_path(str(base_pkg_init)) is not None:
        return str(base_pkg_init)

    base_file = Path(*base_mod.split('.')).with_suffix('.py')
    if fs.get_node_by_path(str(base_file)) is not None:
        return str(base_file)

    return None


def experiment():
    fs = FileSystem.from_disk(Path('data/code4'))

    py_dep = build_py_dep_graph(fs)

    for node in py_dep.nodes:
        print(node)
        print(f'imports: {node.imports}')
        print(f'relative_path: {node.relative_path}')
        print('')
        print(f'node.content: {node.content}')

    # print (py_dep.nodes)

    for e in py_dep.nx_graph.edges():
        print(e)


if __name__ == '__main__':
    # import pytest
    experiment()
    # pytest.main([__file__])
