"""Documentation generator for automatic docs creation.

This module generates documentation in various formats including
docstrings, README files, and API documentation.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Optional, TYPE_CHECKING

from ieum.docs.extractor import (
    CodeExtractor,
    ModuleDoc,
    ClassDoc,
    FunctionDoc,
    ParameterDoc,
)

if TYPE_CHECKING:
    from langchain_core.language_models import BaseChatModel


class DocStyle(Enum):
    """Documentation style formats."""

    GOOGLE = "google"
    NUMPY = "numpy"
    SPHINX = "sphinx"
    EPYTEXT = "epytext"


class DocGenerator:
    """Generates documentation from code analysis."""

    def __init__(self, model: Optional["BaseChatModel"] = None):
        """Initialize the generator.

        Args:
            model: Optional LLM for enhanced doc generation.
        """
        self.model = model
        self.extractor = CodeExtractor()

    def generate_docstring(
        self,
        func: FunctionDoc,
        style: DocStyle = DocStyle.GOOGLE,
    ) -> str:
        """Generate a docstring for a function.

        Args:
            func: Function documentation.
            style: Docstring style.

        Returns:
            Generated docstring.
        """
        if style == DocStyle.GOOGLE:
            return self._generate_google_docstring(func)
        elif style == DocStyle.NUMPY:
            return self._generate_numpy_docstring(func)
        elif style == DocStyle.SPHINX:
            return self._generate_sphinx_docstring(func)
        else:
            return self._generate_google_docstring(func)

    def _generate_google_docstring(self, func: FunctionDoc) -> str:
        """Generate Google-style docstring."""
        lines = ['"""']

        # Description
        desc = func.description or f"{func.name} function."
        lines.append(desc)
        lines.append("")

        # Parameters
        if func.parameters:
            lines.append("Args:")
            for param in func.parameters:
                type_str = f" ({param.type_hint})" if param.type_hint else ""
                desc_str = param.description or "Parameter description."
                default_str = f" Defaults to {param.default}." if param.default else ""
                lines.append(f"    {param.name}{type_str}: {desc_str}{default_str}")
            lines.append("")

        # Returns
        if func.returns:
            lines.append("Returns:")
            type_str = func.returns.type_hint or "value"
            desc_str = func.returns.description or "Return value description."
            lines.append(f"    {type_str}: {desc_str}")
            lines.append("")

        # Raises
        if func.raises:
            lines.append("Raises:")
            for exc in func.raises:
                desc_str = exc.description or "Exception description."
                lines.append(f"    {exc.type_name}: {desc_str}")
            lines.append("")

        # Examples
        if func.examples:
            lines.append("Examples:")
            for example in func.examples:
                lines.append(f"    {example}")
            lines.append("")

        lines.append('"""')
        return "\n".join(lines)

    def _generate_numpy_docstring(self, func: FunctionDoc) -> str:
        """Generate NumPy-style docstring."""
        lines = ['"""']

        # Description
        desc = func.description or f"{func.name} function."
        lines.append(desc)
        lines.append("")

        # Parameters
        if func.parameters:
            lines.append("Parameters")
            lines.append("----------")
            for param in func.parameters:
                type_str = param.type_hint or "type"
                lines.append(f"{param.name} : {type_str}")
                desc_str = param.description or "Parameter description."
                lines.append(f"    {desc_str}")
            lines.append("")

        # Returns
        if func.returns:
            lines.append("Returns")
            lines.append("-------")
            type_str = func.returns.type_hint or "type"
            lines.append(type_str)
            desc_str = func.returns.description or "Return value description."
            lines.append(f"    {desc_str}")
            lines.append("")

        # Raises
        if func.raises:
            lines.append("Raises")
            lines.append("------")
            for exc in func.raises:
                lines.append(exc.type_name)
                desc_str = exc.description or "Exception description."
                lines.append(f"    {desc_str}")
            lines.append("")

        lines.append('"""')
        return "\n".join(lines)

    def _generate_sphinx_docstring(self, func: FunctionDoc) -> str:
        """Generate Sphinx/reST-style docstring."""
        lines = ['"""']

        # Description
        desc = func.description or f"{func.name} function."
        lines.append(desc)
        lines.append("")

        # Parameters
        for param in func.parameters:
            type_str = param.type_hint or "type"
            desc_str = param.description or "Parameter description."
            lines.append(f":param {param.name}: {desc_str}")
            lines.append(f":type {param.name}: {type_str}")

        # Returns
        if func.returns:
            desc_str = func.returns.description or "Return value description."
            lines.append(f":returns: {desc_str}")
            if func.returns.type_hint:
                lines.append(f":rtype: {func.returns.type_hint}")

        # Raises
        for exc in func.raises:
            desc_str = exc.description or "Exception description."
            lines.append(f":raises {exc.type_name}: {desc_str}")

        lines.append('"""')
        return "\n".join(lines)

    def generate_class_docstring(
        self,
        cls: ClassDoc,
        style: DocStyle = DocStyle.GOOGLE,
    ) -> str:
        """Generate a docstring for a class.

        Args:
            cls: Class documentation.
            style: Docstring style.

        Returns:
            Generated docstring.
        """
        lines = ['"""']

        # Description
        desc = cls.description or f"{cls.name} class."
        lines.append(desc)
        lines.append("")

        # Attributes
        if cls.attributes or cls.class_attributes:
            if style == DocStyle.GOOGLE:
                lines.append("Attributes:")
                for attr in cls.class_attributes + cls.attributes:
                    type_str = f" ({attr.type_hint})" if attr.type_hint else ""
                    desc_str = attr.description or "Attribute description."
                    lines.append(f"    {attr.name}{type_str}: {desc_str}")
            elif style == DocStyle.NUMPY:
                lines.append("Attributes")
                lines.append("----------")
                for attr in cls.class_attributes + cls.attributes:
                    type_str = attr.type_hint or "type"
                    lines.append(f"{attr.name} : {type_str}")
                    desc_str = attr.description or "Attribute description."
                    lines.append(f"    {desc_str}")

        lines.append('"""')
        return "\n".join(lines)

    def generate_readme(
        self,
        project_path: str | Path,
        include_api: bool = True,
        template: Optional[str] = None,
    ) -> str:
        """Generate a README.md file for a project.

        Args:
            project_path: Path to the project root.
            include_api: Include API documentation section.
            template: Optional template to use.

        Returns:
            Generated README content.
        """
        project_path = Path(project_path)
        project_name = project_path.name

        sections = []

        # Header
        sections.append(f"# {project_name}\n")

        # Try to get description from existing README or pyproject.toml
        description = self._get_project_description(project_path)
        if description:
            sections.append(f"{description}\n")

        # Installation
        sections.append("## Installation\n")
        if (project_path / "pyproject.toml").exists():
            sections.append("```bash")
            sections.append(f"pip install {project_name}")
            sections.append("```\n")
        elif (project_path / "package.json").exists():
            sections.append("```bash")
            sections.append(f"npm install {project_name}")
            sections.append("```\n")

        # Usage
        sections.append("## Usage\n")
        sections.append("```python")
        sections.append(f"import {project_name.replace('-', '_')}")
        sections.append("```\n")

        # API Documentation
        if include_api:
            api_section = self._generate_api_section(project_path)
            if api_section:
                sections.append("## API Reference\n")
                sections.append(api_section)

        # License
        if (project_path / "LICENSE").exists():
            sections.append("## License\n")
            sections.append("See [LICENSE](LICENSE) file.\n")

        return "\n".join(sections)

    def _get_project_description(self, project_path: Path) -> str:
        """Get project description from metadata files."""
        # Try pyproject.toml
        pyproject = project_path / "pyproject.toml"
        if pyproject.exists():
            try:
                import tomllib
                with open(pyproject, "rb") as f:
                    data = tomllib.load(f)
                    return data.get("project", {}).get("description", "")
            except Exception:
                pass

        # Try existing README
        for readme in ["README.md", "README.rst", "README.txt"]:
            readme_path = project_path / readme
            if readme_path.exists():
                try:
                    content = readme_path.read_text()
                    # Get first paragraph
                    lines = content.split("\n\n")
                    for line in lines:
                        if not line.startswith("#") and line.strip():
                            return line.strip()
                except Exception:
                    pass

        return ""

    def _generate_api_section(self, project_path: Path) -> str:
        """Generate API documentation section."""
        lines = []

        # Find Python files
        py_files = list(project_path.rglob("*.py"))
        py_files = [f for f in py_files if "__pycache__" not in str(f)]

        if not py_files:
            return ""

        for py_file in sorted(py_files)[:10]:  # Limit to 10 files
            try:
                module = self.extractor.extract_module(py_file)

                if module.classes or module.functions:
                    rel_path = py_file.relative_to(project_path)
                    lines.append(f"### `{rel_path}`\n")

                    for cls in module.classes:
                        lines.append(f"#### `class {cls.name}`\n")
                        if cls.description:
                            lines.append(f"{cls.description}\n")

                        for method in cls.methods[:5]:  # Limit methods
                            if not method.name.startswith("_"):
                                lines.append(f"- `{method.signature}`")

                    for func in module.functions[:5]:  # Limit functions
                        if not func.name.startswith("_"):
                            lines.append(f"- `{func.signature}`")

                    lines.append("")

            except Exception:
                continue

        return "\n".join(lines)

    def generate_api_docs(
        self,
        module_path: str | Path,
        format: str = "markdown",
    ) -> str:
        """Generate API documentation for a module.

        Args:
            module_path: Path to the module.
            format: Output format (markdown, rst).

        Returns:
            Generated API documentation.
        """
        module = self.extractor.extract_module(module_path)

        if format == "rst":
            return self._generate_rst_api(module)
        else:
            return self._generate_markdown_api(module)

    def _generate_markdown_api(self, module: ModuleDoc) -> str:
        """Generate Markdown API documentation."""
        lines = []

        lines.append(f"# {module.name}\n")

        if module.description:
            lines.append(f"{module.description}\n")

        # Classes
        for cls in module.classes:
            lines.append(f"## class `{cls.name}`\n")

            if cls.bases:
                lines.append(f"**Bases:** {', '.join(cls.bases)}\n")

            if cls.description:
                lines.append(f"{cls.description}\n")

            # Methods
            for method in cls.methods:
                if method.name.startswith("__") and method.name != "__init__":
                    continue

                lines.append(f"### `{method.signature}`\n")

                if method.description:
                    lines.append(f"{method.description}\n")

                if method.parameters:
                    lines.append("**Parameters:**\n")
                    for param in method.parameters:
                        type_str = f" *({param.type_hint})*" if param.type_hint else ""
                        desc = param.description or ""
                        lines.append(f"- `{param.name}`{type_str}: {desc}")
                    lines.append("")

                if method.returns and method.returns.type_hint:
                    lines.append(f"**Returns:** `{method.returns.type_hint}`")
                    if method.returns.description:
                        lines.append(f"  {method.returns.description}")
                    lines.append("")

        # Functions
        if module.functions:
            lines.append("## Functions\n")

            for func in module.functions:
                if func.name.startswith("_"):
                    continue

                lines.append(f"### `{func.signature}`\n")

                if func.description:
                    lines.append(f"{func.description}\n")

                if func.parameters:
                    lines.append("**Parameters:**\n")
                    for param in func.parameters:
                        type_str = f" *({param.type_hint})*" if param.type_hint else ""
                        desc = param.description or ""
                        lines.append(f"- `{param.name}`{type_str}: {desc}")
                    lines.append("")

                if func.returns and func.returns.type_hint:
                    lines.append(f"**Returns:** `{func.returns.type_hint}`\n")

        return "\n".join(lines)

    def _generate_rst_api(self, module: ModuleDoc) -> str:
        """Generate reStructuredText API documentation."""
        lines = []

        title = module.name
        lines.append(title)
        lines.append("=" * len(title))
        lines.append("")

        if module.description:
            lines.append(module.description)
            lines.append("")

        for cls in module.classes:
            lines.append(f".. py:class:: {cls.name}")
            lines.append("")
            if cls.description:
                lines.append(f"   {cls.description}")
                lines.append("")

            for method in cls.methods:
                if method.name.startswith("__") and method.name != "__init__":
                    continue
                lines.append(f"   .. py:method:: {method.name}({', '.join(p.name for p in method.parameters)})")
                lines.append("")
                if method.description:
                    lines.append(f"      {method.description}")
                    lines.append("")

        return "\n".join(lines)

    async def generate_ai_docstring(
        self,
        func_code: str,
        style: DocStyle = DocStyle.GOOGLE,
    ) -> str:
        """Generate docstring using AI.

        Args:
            func_code: Function source code.
            style: Docstring style.

        Returns:
            AI-generated docstring.
        """
        if not self.model:
            return '"""Function description."""'

        prompt = f"""Generate a {style.value}-style Python docstring for this function.
Only output the docstring, nothing else.

Function:
```python
{func_code}
```

Docstring:"""

        response = self.model.invoke(prompt)
        content = response.content if hasattr(response, "content") else str(response)

        # Clean up response
        content = content.strip()
        if not content.startswith('"""'):
            content = '"""' + content
        if not content.endswith('"""'):
            content = content + '"""'

        return content


# Convenience functions
def generate_docstring(
    file_path: str,
    function_name: str,
    style: str = "google",
) -> str:
    """Generate docstring for a function.

    Args:
        file_path: Path to the file.
        function_name: Name of the function.
        style: Docstring style.

    Returns:
        Generated docstring.
    """
    extractor = CodeExtractor()
    module = extractor.extract_module(file_path)

    # Find function
    for func in module.functions:
        if func.name == function_name:
            generator = DocGenerator()
            return generator.generate_docstring(func, DocStyle(style))

    # Search in classes
    for cls in module.classes:
        for method in cls.methods:
            if method.name == function_name:
                generator = DocGenerator()
                return generator.generate_docstring(method, DocStyle(style))

    return '"""Function not found."""'


def generate_readme(
    project_path: str,
    include_api: bool = True,
) -> str:
    """Generate README for a project.

    Args:
        project_path: Project path.
        include_api: Include API section.

    Returns:
        Generated README content.
    """
    generator = DocGenerator()
    return generator.generate_readme(project_path, include_api)


def generate_api_docs(
    module_path: str,
    format: str = "markdown",
) -> str:
    """Generate API documentation.

    Args:
        module_path: Path to module.
        format: Output format.

    Returns:
        Generated documentation.
    """
    generator = DocGenerator()
    return generator.generate_api_docs(module_path, format)
