"""Code information extractor for documentation generation.

This module extracts documentation-relevant information from code
using AST parsing.
"""

from __future__ import annotations

import ast
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Any

from ieum.ast.base import get_parser_for_file


@dataclass
class ParameterDoc:
    """Documentation for a function parameter."""

    name: str
    type_hint: Optional[str] = None
    default: Optional[str] = None
    description: str = ""


@dataclass
class ReturnDoc:
    """Documentation for a return value."""

    type_hint: Optional[str] = None
    description: str = ""


@dataclass
class ExceptionDoc:
    """Documentation for a raised exception."""

    type_name: str
    description: str = ""


@dataclass
class FunctionDoc:
    """Documentation for a function or method."""

    name: str
    signature: str
    description: str = ""
    parameters: list[ParameterDoc] = field(default_factory=list)
    returns: Optional[ReturnDoc] = None
    raises: list[ExceptionDoc] = field(default_factory=list)
    examples: list[str] = field(default_factory=list)
    source_file: str = ""
    line: int = 0
    is_async: bool = False
    is_method: bool = False
    decorators: list[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return {
            "name": self.name,
            "signature": self.signature,
            "description": self.description,
            "parameters": [
                {"name": p.name, "type": p.type_hint, "default": p.default, "description": p.description}
                for p in self.parameters
            ],
            "returns": {"type": self.returns.type_hint, "description": self.returns.description} if self.returns else None,
            "raises": [{"type": e.type_name, "description": e.description} for e in self.raises],
            "examples": self.examples,
            "source_file": self.source_file,
            "line": self.line,
            "is_async": self.is_async,
            "is_method": self.is_method,
            "decorators": self.decorators,
        }


@dataclass
class AttributeDoc:
    """Documentation for a class attribute."""

    name: str
    type_hint: Optional[str] = None
    description: str = ""
    default: Optional[str] = None


@dataclass
class ClassDoc:
    """Documentation for a class."""

    name: str
    description: str = ""
    bases: list[str] = field(default_factory=list)
    methods: list[FunctionDoc] = field(default_factory=list)
    attributes: list[AttributeDoc] = field(default_factory=list)
    class_attributes: list[AttributeDoc] = field(default_factory=list)
    source_file: str = ""
    line: int = 0
    decorators: list[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return {
            "name": self.name,
            "description": self.description,
            "bases": self.bases,
            "methods": [m.to_dict() for m in self.methods],
            "attributes": [{"name": a.name, "type": a.type_hint, "description": a.description} for a in self.attributes],
            "class_attributes": [{"name": a.name, "type": a.type_hint, "description": a.description} for a in self.class_attributes],
            "source_file": self.source_file,
            "line": self.line,
            "decorators": self.decorators,
        }


@dataclass
class ModuleDoc:
    """Documentation for a module."""

    name: str
    path: str
    description: str = ""
    classes: list[ClassDoc] = field(default_factory=list)
    functions: list[FunctionDoc] = field(default_factory=list)
    constants: list[AttributeDoc] = field(default_factory=list)
    imports: list[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return {
            "name": self.name,
            "path": self.path,
            "description": self.description,
            "classes": [c.to_dict() for c in self.classes],
            "functions": [f.to_dict() for f in self.functions],
            "constants": [{"name": c.name, "type": c.type_hint, "description": c.description} for c in self.constants],
            "imports": self.imports,
        }


class CodeExtractor:
    """Extracts documentation information from code."""

    def __init__(self):
        """Initialize the extractor."""
        pass

    def extract_module(self, file_path: str | Path) -> ModuleDoc:
        """Extract documentation from a module.

        Args:
            file_path: Path to the Python file.

        Returns:
            ModuleDoc with extracted information.
        """
        file_path = Path(file_path)
        module_name = file_path.stem

        module_doc = ModuleDoc(
            name=module_name,
            path=str(file_path),
        )

        try:
            with open(file_path, "r", encoding="utf-8") as f:
                source = f.read()

            tree = ast.parse(source)

            # Get module docstring
            module_doc.description = ast.get_docstring(tree) or ""

            # Extract imports
            for node in ast.walk(tree):
                if isinstance(node, ast.Import):
                    for alias in node.names:
                        module_doc.imports.append(alias.name)
                elif isinstance(node, ast.ImportFrom):
                    module_name_import = node.module or ""
                    for alias in node.names:
                        module_doc.imports.append(f"{module_name_import}.{alias.name}")

            # Extract top-level elements
            for node in ast.iter_child_nodes(tree):
                if isinstance(node, ast.ClassDef):
                    class_doc = self._extract_class(node, str(file_path))
                    module_doc.classes.append(class_doc)

                elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    func_doc = self._extract_function(node, str(file_path))
                    module_doc.functions.append(func_doc)

                elif isinstance(node, ast.Assign):
                    # Module-level constants
                    for target in node.targets:
                        if isinstance(target, ast.Name):
                            if target.id.isupper():  # Likely a constant
                                module_doc.constants.append(AttributeDoc(
                                    name=target.id,
                                    type_hint=self._get_type_annotation(node.value),
                                ))

        except Exception as e:
            module_doc.description = f"Error parsing module: {e}"

        return module_doc

    def _extract_class(self, node: ast.ClassDef, file_path: str) -> ClassDoc:
        """Extract documentation from a class."""
        class_doc = ClassDoc(
            name=node.name,
            description=ast.get_docstring(node) or "",
            source_file=file_path,
            line=node.lineno,
        )

        # Base classes
        for base in node.bases:
            try:
                class_doc.bases.append(ast.unparse(base))
            except Exception:
                pass

        # Decorators
        for decorator in node.decorator_list:
            try:
                class_doc.decorators.append(ast.unparse(decorator))
            except Exception:
                pass

        # Class body
        for child in node.body:
            if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
                method_doc = self._extract_function(child, file_path, is_method=True)
                class_doc.methods.append(method_doc)

            elif isinstance(child, ast.AnnAssign):
                # Type-annotated class attribute
                if isinstance(child.target, ast.Name):
                    attr = AttributeDoc(
                        name=child.target.id,
                        type_hint=ast.unparse(child.annotation) if child.annotation else None,
                    )
                    class_doc.class_attributes.append(attr)

            elif isinstance(child, ast.Assign):
                # Class attribute without annotation
                for target in child.targets:
                    if isinstance(target, ast.Name):
                        class_doc.class_attributes.append(AttributeDoc(
                            name=target.id,
                        ))

        return class_doc

    def _extract_function(
        self,
        node: ast.FunctionDef | ast.AsyncFunctionDef,
        file_path: str,
        is_method: bool = False,
    ) -> FunctionDoc:
        """Extract documentation from a function."""
        func_doc = FunctionDoc(
            name=node.name,
            signature=self._build_signature(node),
            description="",
            source_file=file_path,
            line=node.lineno,
            is_async=isinstance(node, ast.AsyncFunctionDef),
            is_method=is_method,
        )

        # Get docstring and parse it
        docstring = ast.get_docstring(node)
        if docstring:
            parsed = self._parse_docstring(docstring)
            func_doc.description = parsed.get("description", "")
            func_doc.raises = parsed.get("raises", [])
            func_doc.examples = parsed.get("examples", [])

            # Merge parsed parameters with signature parameters
            param_docs = {p["name"]: p for p in parsed.get("parameters", [])}
            for param in func_doc.parameters:
                if param.name in param_docs:
                    param.description = param_docs[param.name].get("description", "")

            if parsed.get("returns"):
                func_doc.returns = ReturnDoc(
                    type_hint=func_doc.returns.type_hint if func_doc.returns else None,
                    description=parsed["returns"].get("description", ""),
                )

        # Extract parameters from signature
        func_doc.parameters = self._extract_parameters(node.args)

        # Return type
        if node.returns:
            try:
                return_type = ast.unparse(node.returns)
                if func_doc.returns:
                    func_doc.returns.type_hint = return_type
                else:
                    func_doc.returns = ReturnDoc(type_hint=return_type)
            except Exception:
                pass

        # Decorators
        for decorator in node.decorator_list:
            try:
                func_doc.decorators.append(ast.unparse(decorator))
            except Exception:
                pass

        return func_doc

    def _extract_parameters(self, args: ast.arguments) -> list[ParameterDoc]:
        """Extract parameters from function arguments."""
        parameters = []

        # Calculate defaults offset
        defaults_offset = len(args.args) - len(args.defaults)

        for i, arg in enumerate(args.args):
            # Skip 'self' and 'cls'
            if arg.arg in ("self", "cls"):
                continue

            param = ParameterDoc(name=arg.arg)

            # Type hint
            if arg.annotation:
                try:
                    param.type_hint = ast.unparse(arg.annotation)
                except Exception:
                    pass

            # Default value
            if i >= defaults_offset:
                try:
                    param.default = ast.unparse(args.defaults[i - defaults_offset])
                except Exception:
                    pass

            parameters.append(param)

        # *args
        if args.vararg:
            param = ParameterDoc(name=f"*{args.vararg.arg}")
            if args.vararg.annotation:
                try:
                    param.type_hint = ast.unparse(args.vararg.annotation)
                except Exception:
                    pass
            parameters.append(param)

        # **kwargs
        if args.kwarg:
            param = ParameterDoc(name=f"**{args.kwarg.arg}")
            if args.kwarg.annotation:
                try:
                    param.type_hint = ast.unparse(args.kwarg.annotation)
                except Exception:
                    pass
            parameters.append(param)

        return parameters

    def _build_signature(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> str:
        """Build function signature string."""
        try:
            parts = []

            # Parameters
            param_strs = []
            for param in self._extract_parameters(node.args):
                p = param.name
                if param.type_hint:
                    p += f": {param.type_hint}"
                if param.default:
                    p += f" = {param.default}"
                param_strs.append(p)

            params = ", ".join(param_strs)

            # Async prefix
            prefix = "async " if isinstance(node, ast.AsyncFunctionDef) else ""

            sig = f"{prefix}def {node.name}({params})"

            # Return type
            if node.returns:
                sig += f" -> {ast.unparse(node.returns)}"

            return sig

        except Exception:
            return f"def {node.name}(...)"

    def _parse_docstring(self, docstring: str) -> dict:
        """Parse docstring to extract structured information.

        Supports Google and NumPy style docstrings.
        """
        result = {
            "description": "",
            "parameters": [],
            "returns": None,
            "raises": [],
            "examples": [],
        }

        if not docstring:
            return result

        lines = docstring.strip().split("\n")

        # First paragraph is description
        desc_lines = []
        i = 0
        while i < len(lines) and lines[i].strip() and not self._is_section_header(lines[i]):
            desc_lines.append(lines[i])
            i += 1

        result["description"] = " ".join(desc_lines).strip()

        # Parse sections
        current_section = None
        section_content = []

        while i < len(lines):
            line = lines[i]

            if self._is_section_header(line):
                # Process previous section
                if current_section:
                    self._process_section(current_section, section_content, result)

                current_section = line.strip().rstrip(":").lower()
                section_content = []
            else:
                section_content.append(line)

            i += 1

        # Process last section
        if current_section:
            self._process_section(current_section, section_content, result)

        return result

    def _is_section_header(self, line: str) -> bool:
        """Check if line is a section header."""
        headers = ["args", "arguments", "parameters", "params",
                   "returns", "return", "yields", "yield",
                   "raises", "raise", "exceptions",
                   "examples", "example", "note", "notes"]

        stripped = line.strip().rstrip(":").lower()
        return stripped in headers

    def _process_section(self, section: str, content: list[str], result: dict) -> None:
        """Process a docstring section."""
        if section in ("args", "arguments", "parameters", "params"):
            for line in content:
                param = self._parse_param_line(line)
                if param:
                    result["parameters"].append(param)

        elif section in ("returns", "return", "yields", "yield"):
            description = " ".join(line.strip() for line in content if line.strip())
            result["returns"] = {"description": description}

        elif section in ("raises", "raise", "exceptions"):
            for line in content:
                exc = self._parse_exception_line(line)
                if exc:
                    result["raises"].append(exc)

        elif section in ("examples", "example"):
            result["examples"] = [line for line in content if line.strip()]

    def _parse_param_line(self, line: str) -> Optional[dict]:
        """Parse a parameter documentation line."""
        line = line.strip()
        if not line:
            return None

        # Google style: name (type): description
        # NumPy style: name : type\n    description

        match = re.match(r"(\w+)\s*(?:\(([^)]+)\))?\s*:?\s*(.*)", line)
        if match:
            return {
                "name": match.group(1),
                "type": match.group(2),
                "description": match.group(3).strip(),
            }

        return None

    def _parse_exception_line(self, line: str) -> Optional[ExceptionDoc]:
        """Parse an exception documentation line."""
        line = line.strip()
        if not line:
            return None

        match = re.match(r"(\w+)\s*:?\s*(.*)", line)
        if match:
            return ExceptionDoc(
                type_name=match.group(1),
                description=match.group(2).strip(),
            )

        return None

    def _get_type_annotation(self, node: ast.AST) -> Optional[str]:
        """Get type annotation from an AST node."""
        try:
            if isinstance(node, ast.Constant):
                return type(node.value).__name__
            elif isinstance(node, ast.List):
                return "list"
            elif isinstance(node, ast.Dict):
                return "dict"
            elif isinstance(node, ast.Set):
                return "set"
            elif isinstance(node, ast.Tuple):
                return "tuple"
        except Exception:
            pass
        return None


# Convenience functions
def extract_module_doc(file_path: str | Path) -> dict:
    """Extract documentation from a module.

    Args:
        file_path: Path to the file.

    Returns:
        Module documentation dictionary.
    """
    extractor = CodeExtractor()
    doc = extractor.extract_module(file_path)
    return doc.to_dict()


def extract_function_docs(file_path: str | Path) -> list[dict]:
    """Extract function documentation from a file.

    Args:
        file_path: Path to the file.

    Returns:
        List of function documentation dictionaries.
    """
    extractor = CodeExtractor()
    module = extractor.extract_module(file_path)

    docs = [f.to_dict() for f in module.functions]
    for cls in module.classes:
        for method in cls.methods:
            method_dict = method.to_dict()
            method_dict["class"] = cls.name
            docs.append(method_dict)

    return docs
