#!/usr/bin/env python3
"""
GENESIS API DOCUMENTATION GENERATOR
=====================================
Auto-generates API documentation from Python modules.

Features:
    - Docstring extraction
    - Type hint documentation
    - Class hierarchy
    - Function signatures
    - Example extraction
    - Markdown output

Usage:
    generator = APIDocGenerator()
    docs = generator.generate_for_module("memory_integration")
    generator.generate_all("core", output_dir="docs/api")
"""

import ast
import inspect
import importlib
import importlib.util
import json
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime


@dataclass
class ParamDoc:
    """Documentation for a parameter."""
    name: str
    type_hint: str = ""
    description: str = ""
    default: str = ""

    def to_dict(self) -> Dict:
        return {
            "name": self.name,
            "type": self.type_hint,
            "description": self.description,
            "default": self.default
        }


@dataclass
class FunctionDoc:
    """Documentation for a function."""
    name: str
    signature: str
    docstring: str = ""
    params: List[ParamDoc] = field(default_factory=list)
    returns: str = ""
    return_type: str = ""
    raises: List[str] = field(default_factory=list)
    examples: List[str] = field(default_factory=list)
    is_async: bool = False
    is_classmethod: bool = False
    is_staticmethod: bool = False

    def to_dict(self) -> Dict:
        return {
            "name": self.name,
            "signature": self.signature,
            "docstring": self.docstring,
            "params": [p.to_dict() for p in self.params],
            "returns": self.returns,
            "return_type": self.return_type,
            "raises": self.raises,
            "examples": self.examples,
            "is_async": self.is_async
        }


@dataclass
class ClassDoc:
    """Documentation for a class."""
    name: str
    docstring: str = ""
    bases: List[str] = field(default_factory=list)
    methods: List[FunctionDoc] = field(default_factory=list)
    attributes: List[ParamDoc] = field(default_factory=list)
    class_methods: List[FunctionDoc] = field(default_factory=list)
    static_methods: List[FunctionDoc] = field(default_factory=list)

    def to_dict(self) -> Dict:
        return {
            "name": self.name,
            "docstring": self.docstring,
            "bases": self.bases,
            "methods": [m.to_dict() for m in self.methods],
            "attributes": [a.to_dict() for a in self.attributes],
            "class_methods": [m.to_dict() for m in self.class_methods],
            "static_methods": [m.to_dict() for m in self.static_methods]
        }


@dataclass
class ModuleDoc:
    """Documentation for a module."""
    name: str
    path: str
    docstring: str = ""
    classes: List[ClassDoc] = field(default_factory=list)
    functions: List[FunctionDoc] = field(default_factory=list)
    constants: List[Tuple[str, Any]] = field(default_factory=list)

    def to_dict(self) -> Dict:
        return {
            "name": self.name,
            "path": self.path,
            "docstring": self.docstring,
            "classes": [c.to_dict() for c in self.classes],
            "functions": [f.to_dict() for f in self.functions],
            "constants": [(n, str(v)) for n, v in self.constants]
        }


class DocstringParser:
    """
    Parses docstrings in Google/NumPy style.
    """

    @staticmethod
    def parse(docstring: str) -> Dict:
        """Parse a docstring into sections."""
        if not docstring:
            return {"description": "", "params": [], "returns": "", "raises": [], "examples": []}

        sections = {
            "description": "",
            "params": [],
            "returns": "",
            "raises": [],
            "examples": []
        }

        lines = docstring.strip().split('\n')
        current_section = "description"
        current_param = None
        buffer = []

        section_headers = {
            "Args:": "params",
            "Arguments:": "params",
            "Parameters:": "params",
            "Returns:": "returns",
            "Return:": "returns",
            "Raises:": "raises",
            "Exceptions:": "raises",
            "Examples:": "examples",
            "Example:": "examples"
        }

        for line in lines:
            stripped = line.strip()

            # Check for section header
            if stripped in section_headers:
                # Save buffer
                if buffer and current_section == "description":
                    sections["description"] = '\n'.join(buffer).strip()
                buffer = []
                current_section = section_headers[stripped]
                current_param = None
                continue

            # Handle content based on current section
            if current_section == "params":
                # Look for parameter definition: "name (type): description"
                param_match = re.match(r'^\s*(\w+)\s*(?:\(([^)]+)\))?\s*:\s*(.*)$', line)
                if param_match:
                    if current_param:
                        sections["params"].append(current_param)
                    current_param = {
                        "name": param_match.group(1),
                        "type": param_match.group(2) or "",
                        "description": param_match.group(3)
                    }
                elif current_param and stripped:
                    current_param["description"] += " " + stripped

            elif current_section == "returns":
                if stripped:
                    if sections["returns"]:
                        sections["returns"] += " " + stripped
                    else:
                        sections["returns"] = stripped

            elif current_section == "raises":
                if stripped:
                    sections["raises"].append(stripped)

            elif current_section == "examples":
                buffer.append(line)

            else:  # description
                buffer.append(line)

        # Finalize
        if current_param:
            sections["params"].append(current_param)

        if buffer:
            if current_section == "description":
                sections["description"] = '\n'.join(buffer).strip()
            elif current_section == "examples":
                sections["examples"] = ['\n'.join(buffer)]

        return sections


class ASTDocExtractor:
    """
    Extracts documentation using AST.
    """

    def __init__(self, source_code: str):
        self.source = source_code
        self.tree = ast.parse(source_code)
        self.parser = DocstringParser()

    def get_type_annotation(self, annotation: ast.expr) -> str:
        """Convert AST annotation to string."""
        if annotation is None:
            return ""

        if isinstance(annotation, ast.Name):
            return annotation.id
        elif isinstance(annotation, ast.Constant):
            return str(annotation.value)
        elif isinstance(annotation, ast.Subscript):
            base = self.get_type_annotation(annotation.value)
            if isinstance(annotation.slice, ast.Tuple):
                args = ', '.join(self.get_type_annotation(e) for e in annotation.slice.elts)
            else:
                args = self.get_type_annotation(annotation.slice)
            return f"{base}[{args}]"
        elif isinstance(annotation, ast.Attribute):
            return f"{self.get_type_annotation(annotation.value)}.{annotation.attr}"
        elif isinstance(annotation, ast.BinOp) and isinstance(annotation.op, ast.BitOr):
            # Union types: X | Y
            left = self.get_type_annotation(annotation.left)
            right = self.get_type_annotation(annotation.right)
            return f"{left} | {right}"
        else:
            return ast.unparse(annotation) if hasattr(ast, 'unparse') else "Unknown"

    def extract_function(self, node: ast.FunctionDef) -> FunctionDoc:
        """Extract documentation for a function."""
        # Get signature
        args = []
        for arg in node.args.args:
            arg_str = arg.arg
            if arg.annotation:
                arg_str += f": {self.get_type_annotation(arg.annotation)}"
            args.append(arg_str)

        # Handle defaults
        defaults = node.args.defaults
        if defaults:
            for i, default in enumerate(defaults):
                idx = len(args) - len(defaults) + i
                if isinstance(default, ast.Constant):
                    args[idx] += f" = {repr(default.value)}"
                else:
                    args[idx] += " = ..."

        signature = f"({', '.join(args)})"

        # Return type
        return_type = ""
        if node.returns:
            return_type = self.get_type_annotation(node.returns)
            signature += f" -> {return_type}"

        # Docstring
        docstring = ast.get_docstring(node) or ""
        parsed = self.parser.parse(docstring)

        # Build params
        params = []
        for param in parsed["params"]:
            params.append(ParamDoc(
                name=param["name"],
                type_hint=param["type"],
                description=param["description"]
            ))

        # Check decorators
        is_async = isinstance(node, ast.AsyncFunctionDef)
        is_classmethod = any(
            isinstance(d, ast.Name) and d.id == "classmethod"
            for d in node.decorator_list
        )
        is_staticmethod = any(
            isinstance(d, ast.Name) and d.id == "staticmethod"
            for d in node.decorator_list
        )

        return FunctionDoc(
            name=node.name,
            signature=signature,
            docstring=parsed["description"],
            params=params,
            returns=parsed["returns"],
            return_type=return_type,
            raises=parsed["raises"],
            examples=parsed["examples"],
            is_async=is_async,
            is_classmethod=is_classmethod,
            is_staticmethod=is_staticmethod
        )

    def extract_class(self, node: ast.ClassDef) -> ClassDoc:
        """Extract documentation for a class."""
        # Get base classes
        bases = []
        for base in node.bases:
            if isinstance(base, ast.Name):
                bases.append(base.id)
            elif isinstance(base, ast.Attribute):
                bases.append(f"{self.get_type_annotation(base.value)}.{base.attr}")

        # Docstring
        docstring = ast.get_docstring(node) or ""

        # Extract methods
        methods = []
        class_methods = []
        static_methods = []
        attributes = []

        for item in node.body:
            if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
                func_doc = self.extract_function(item)

                if func_doc.name.startswith('_') and not func_doc.name.startswith('__'):
                    continue  # Skip private methods

                if func_doc.is_classmethod:
                    class_methods.append(func_doc)
                elif func_doc.is_staticmethod:
                    static_methods.append(func_doc)
                else:
                    methods.append(func_doc)

            elif isinstance(item, ast.AnnAssign):
                # Class attribute with annotation
                if isinstance(item.target, ast.Name):
                    attributes.append(ParamDoc(
                        name=item.target.id,
                        type_hint=self.get_type_annotation(item.annotation),
                        default=ast.unparse(item.value) if item.value and hasattr(ast, 'unparse') else ""
                    ))

        return ClassDoc(
            name=node.name,
            docstring=docstring,
            bases=bases,
            methods=methods,
            class_methods=class_methods,
            static_methods=static_methods,
            attributes=attributes
        )

    def extract_module(self, name: str, path: str) -> ModuleDoc:
        """Extract documentation for the entire module."""
        docstring = ast.get_docstring(self.tree) or ""

        classes = []
        functions = []
        constants = []

        for node in self.tree.body:
            if isinstance(node, ast.ClassDef):
                if not node.name.startswith('_'):
                    classes.append(self.extract_class(node))

            elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                if not node.name.startswith('_'):
                    functions.append(self.extract_function(node))

            elif isinstance(node, ast.Assign):
                for target in node.targets:
                    if isinstance(target, ast.Name) and target.id.isupper():
                        if isinstance(node.value, ast.Constant):
                            constants.append((target.id, node.value.value))

        return ModuleDoc(
            name=name,
            path=path,
            docstring=docstring,
            classes=classes,
            functions=functions,
            constants=constants
        )


class APIDocGenerator:
    """
    Generates API documentation for Genesis modules.
    """

    def __init__(self, project_root: Path = None):
        self.project_root = project_root or Path(__file__).parent.parent

    def generate_for_file(self, file_path: Path) -> Optional[ModuleDoc]:
        """Generate documentation for a single file."""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                source = f.read()

            extractor = ASTDocExtractor(source)
            return extractor.extract_module(
                name=file_path.stem,
                path=str(file_path.relative_to(self.project_root))
            )
        except Exception as e:
            print(f"Error processing {file_path}: {e}")
            return None

    def generate_for_directory(self, directory: str = "core") -> List[ModuleDoc]:
        """Generate documentation for all modules in a directory."""
        target_dir = self.project_root / directory
        docs = []

        if not target_dir.exists():
            return docs

        for py_file in sorted(target_dir.glob("*.py")):
            if py_file.name.startswith('_'):
                continue

            doc = self.generate_for_file(py_file)
            if doc:
                docs.append(doc)

        return docs

    def generate_markdown(self, module_doc: ModuleDoc) -> str:
        """Generate Markdown documentation for a module."""
        lines = [
            f"# {module_doc.name}",
            "",
            f"**Path:** `{module_doc.path}`",
            "",
        ]

        if module_doc.docstring:
            lines.extend([module_doc.docstring, ""])

        # Constants
        if module_doc.constants:
            lines.extend(["## Constants", ""])
            for name, value in module_doc.constants:
                lines.append(f"- `{name}` = `{value}`")
            lines.append("")

        # Classes
        for cls in module_doc.classes:
            lines.extend([
                f"## class `{cls.name}`",
                ""
            ])

            if cls.bases:
                lines.append(f"**Inherits from:** {', '.join(cls.bases)}")
                lines.append("")

            if cls.docstring:
                lines.extend([cls.docstring, ""])

            # Attributes
            if cls.attributes:
                lines.extend(["### Attributes", ""])
                for attr in cls.attributes:
                    type_str = f": `{attr.type_hint}`" if attr.type_hint else ""
                    default_str = f" = `{attr.default}`" if attr.default else ""
                    lines.append(f"- **{attr.name}**{type_str}{default_str}")
                lines.append("")

            # Methods
            if cls.methods:
                lines.extend(["### Methods", ""])
                for method in cls.methods:
                    lines.extend(self._format_function(method, indent=4))

            if cls.class_methods:
                lines.extend(["### Class Methods", ""])
                for method in cls.class_methods:
                    lines.extend(self._format_function(method, indent=4))

            if cls.static_methods:
                lines.extend(["### Static Methods", ""])
                for method in cls.static_methods:
                    lines.extend(self._format_function(method, indent=4))

        # Module-level functions
        if module_doc.functions:
            lines.extend(["## Functions", ""])
            for func in module_doc.functions:
                lines.extend(self._format_function(func))

        return '\n'.join(lines)

    def _format_function(self, func: FunctionDoc, indent: int = 0) -> List[str]:
        """Format a function for markdown."""
        prefix = " " * indent
        lines = []

        async_prefix = "async " if func.is_async else ""
        lines.append(f"#### `{async_prefix}{func.name}{func.signature}`")
        lines.append("")

        if func.docstring:
            lines.append(func.docstring)
            lines.append("")

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

        if func.returns:
            lines.append(f"**Returns:** {func.returns}")
            lines.append("")

        if func.raises:
            lines.append("**Raises:**")
            for exc in func.raises:
                lines.append(f"- {exc}")
            lines.append("")

        return lines

    def generate_all(self, directory: str = "core", output_dir: str = "docs/api"):
        """Generate documentation for all modules and save to files."""
        output_path = self.project_root / output_dir
        output_path.mkdir(parents=True, exist_ok=True)

        docs = self.generate_for_directory(directory)

        # Generate index
        index_lines = [
            f"# Genesis API Documentation",
            f"",
            f"**Generated:** {datetime.now().isoformat()}",
            f"**Modules:** {len(docs)}",
            "",
            "## Modules",
            ""
        ]

        for doc in docs:
            index_lines.append(f"- [{doc.name}]({doc.name}.md) - {doc.docstring.split(chr(10))[0] if doc.docstring else 'No description'}")

            # Write individual module doc
            md_content = self.generate_markdown(doc)
            md_path = output_path / f"{doc.name}.md"
            with open(md_path, 'w', encoding='utf-8') as f:
                f.write(md_content)

        # Write index
        index_path = output_path / "README.md"
        with open(index_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(index_lines))

        return len(docs)


def main():
    """CLI for API documentation generator."""
    import argparse
    parser = argparse.ArgumentParser(description="Genesis API Documentation Generator")
    parser.add_argument("command", choices=["generate", "preview", "json"])
    parser.add_argument("--module", help="Specific module to document")
    parser.add_argument("--dir", default="core", help="Directory to document")
    parser.add_argument("--output", default="docs/api", help="Output directory")
    args = parser.parse_args()

    generator = APIDocGenerator()

    if args.command == "generate":
        count = generator.generate_all(args.dir, args.output)
        print(f"Generated documentation for {count} modules in {args.output}/")

    elif args.command == "preview":
        if args.module:
            file_path = generator.project_root / args.dir / f"{args.module}.py"
            if file_path.exists():
                doc = generator.generate_for_file(file_path)
                if doc:
                    print(generator.generate_markdown(doc))
            else:
                print(f"Module not found: {file_path}")
        else:
            print("--module required for preview")

    elif args.command == "json":
        docs = generator.generate_for_directory(args.dir)
        output = [d.to_dict() for d in docs]
        print(json.dumps(output, indent=2))


if __name__ == "__main__":
    main()