#!/usr/bin/env python3
"""
OmniClaw Living Documentation
Auto-generates and updates Mermaid diagrams as the project evolves.
Architecture diagrams that actually stay accurate.
"""

import logging
import ast
import os
import re
import time
from typing import Dict, List, Optional, Any, Set, Tuple
from pathlib import Path
from dataclasses import dataclass, field

logger = logging.getLogger("OmniClaw.LivingDocs")


@dataclass
class ModuleInfo:
    """Information about a Python module"""
    name: str
    path: str
    classes: List[str] = field(default_factory=list)
    functions: List[str] = field(default_factory=list)
    imports: List[str] = field(default_factory=list)
    docstring: str = ""
    line_count: int = 0


@dataclass
class ClassInfo:
    """Information about a class"""
    name: str
    module: str
    bases: List[str] = field(default_factory=list)
    methods: List[str] = field(default_factory=list)
    attributes: List[str] = field(default_factory=list)
    docstring: str = ""


class LivingDocumentation:
    """
    Auto-generates Mermaid diagrams from Python source code.
    
    Scans imports, class hierarchies, and function call graphs
    to generate architecture diagrams that stay in sync with code.
    """
    
    def __init__(self):
        self._module_cache: Dict[str, ModuleInfo] = {}
        self._class_cache: Dict[str, ClassInfo] = {}
        self._last_scan: float = 0
        
        logger.info("LivingDocumentation initialized")
    
    def generate_architecture_diagram(self, root_dir: str,
                                        max_depth: int = 3) -> str:
        """
        Generate a module dependency graph in Mermaid format.
        
        Args:
            root_dir: Project root directory
            max_depth: Maximum directory depth to scan
            
        Returns:
            Mermaid diagram string
        """
        root = Path(root_dir).resolve()
        modules = self._scan_modules(root, max_depth)
        
        lines = ["graph TD"]
        lines.append("    %% Auto-generated by OmniClaw LivingDocumentation")
        lines.append(f"    %% Generated at {time.strftime('%Y-%m-%d %H:%M:%S')}")
        lines.append("")
        
        # Create node definitions
        seen_nodes = set()
        for mod in modules.values():
            node_id = self._sanitize_id(mod.name)
            if node_id not in seen_nodes:
                label = mod.name
                if mod.classes:
                    label += f"\\n({len(mod.classes)} classes)"
                lines.append(f'    {node_id}["{label}"]')
                seen_nodes.add(node_id)
        
        lines.append("")
        
        # Create edges from imports
        seen_edges = set()
        for mod in modules.values():
            from_id = self._sanitize_id(mod.name)
            for imp in mod.imports:
                # Only show internal imports
                imp_name = imp.split('.')[0] if '.' in imp else imp
                to_id = self._sanitize_id(imp_name)
                
                if to_id in seen_nodes and from_id != to_id:
                    edge = (from_id, to_id)
                    if edge not in seen_edges:
                        lines.append(f"    {from_id} --> {to_id}")
                        seen_edges.add(edge)
        
        lines.append("")
        
        # Style nodes by type
        lines.append("    %% Styling")
        for mod in modules.values():
            node_id = self._sanitize_id(mod.name)
            if mod.classes:
                lines.append(f"    style {node_id} fill:#4a9eff,color:#fff")
            elif mod.functions:
                lines.append(f"    style {node_id} fill:#2ec4b6,color:#fff")
            else:
                lines.append(f"    style {node_id} fill:#e0e0e0,color:#333")
        
        diagram = "\n".join(lines)
        logger.info(f"Architecture diagram generated: {len(modules)} modules, {len(seen_edges)} edges")
        return diagram
    
    def generate_class_diagram(self, files: List[str] = None,
                                 root_dir: str = None) -> str:
        """
        Generate a class hierarchy diagram in Mermaid format.
        
        Args:
            files: Specific files to analyze (or None for all)
            root_dir: Project root (used if files is None)
            
        Returns:
            Mermaid class diagram string
        """
        if files is None and root_dir:
            files = [str(p) for p in Path(root_dir).rglob("*.py")]
        elif files is None:
            return "classDiagram\n    %% No files to analyze"
        
        classes = self._extract_classes(files)
        
        lines = ["classDiagram"]
        lines.append("    %% Auto-generated by OmniClaw LivingDocumentation")
        lines.append("")
        
        for cls in classes.values():
            # Class definition
            lines.append(f"    class {cls.name} {{")
            
            # Attributes
            for attr in cls.attributes[:10]:  # Limit to 10 attributes
                lines.append(f"        {attr}")
            
            # Methods
            for method in cls.methods[:15]:  # Limit to 15 methods
                if method.startswith("_"):
                    lines.append(f"        -{method}()")
                else:
                    lines.append(f"        +{method}()")
            
            lines.append("    }")
            lines.append("")
        
        # Inheritance relationships
        for cls in classes.values():
            for base in cls.bases:
                if base in classes:
                    lines.append(f"    {base} <|-- {cls.name}")
        
        # Composition relationships (heuristic: attributes that match class names)
        for cls in classes.values():
            for attr in cls.attributes:
                # Check if attribute type matches a known class
                for other_name in classes:
                    if other_name.lower() in attr.lower() and other_name != cls.name:
                        lines.append(f"    {cls.name} --> {other_name}")
                        break
        
        diagram = "\n".join(lines)
        logger.info(f"Class diagram generated: {len(classes)} classes")
        return diagram
    
    def generate_flow_diagram(self, function_name: str, 
                                file_path: str) -> str:
        """
        Generate a flow diagram for a function's call graph.
        
        Args:
            function_name: Name of the function to trace
            file_path: Path to the file containing the function
            
        Returns:
            Mermaid flowchart string
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                source = f.read()
            
            tree = ast.parse(source)
        except Exception as e:
            return f"graph TD\n    error[\"{function_name}: Parse error - {e}\"]"
        
        # Find the function
        calls = self._extract_calls(tree, function_name)
        
        lines = ["graph TD"]
        lines.append(f"    %% Call graph for {function_name}")
        lines.append("")
        
        root_id = self._sanitize_id(function_name)
        lines.append(f'    {root_id}(("{function_name}"))')
        lines.append(f"    style {root_id} fill:#ff6b6b,color:#fff")
        
        seen = set()
        for i, (caller, callee, context) in enumerate(calls):
            caller_id = self._sanitize_id(caller)
            callee_id = self._sanitize_id(f"{callee}_{i}")
            
            edge = (caller_id, callee)
            if edge not in seen:
                label = callee
                if context:
                    label += f"\\n({context})"
                lines.append(f'    {caller_id} --> {callee_id}["{label}"]')
                seen.add(edge)
        
        if not calls:
            lines.append(f'    {root_id} --> none["No calls found"]')
        
        diagram = "\n".join(lines)
        logger.info(f"Flow diagram generated for {function_name}: {len(calls)} calls")
        return diagram
    
    def update_docs(self, project_root: str, 
                     output_dir: str = None) -> Dict[str, str]:
        """
        Regenerate all diagrams and write to the output directory.
        
        Args:
            project_root: Project root directory
            output_dir: Output directory (defaults to project_root/docs/diagrams)
            
        Returns:
            Dict of diagram_name -> file_path
        """
        if output_dir is None:
            output_dir = os.path.join(project_root, "docs", "diagrams")
        
        os.makedirs(output_dir, exist_ok=True)
        generated = {}
        
        # Architecture diagram
        arch = self.generate_architecture_diagram(project_root)
        arch_path = os.path.join(output_dir, "architecture.md")
        self._write_diagram(arch_path, "Architecture", arch)
        generated["architecture"] = arch_path
        
        # Class diagram
        cls_diagram = self.generate_class_diagram(root_dir=project_root)
        cls_path = os.path.join(output_dir, "classes.md")
        self._write_diagram(cls_path, "Class Hierarchy", cls_diagram)
        generated["classes"] = cls_path
        
        self._last_scan = time.time()
        logger.info(f"Documentation updated: {len(generated)} diagrams in {output_dir}")
        
        return generated
    
    # --- Private helpers ---
    
    def _scan_modules(self, root: Path, max_depth: int) -> Dict[str, ModuleInfo]:
        """Scan all Python modules in a directory"""
        modules = {}
        ignore = {".git", "__pycache__", ".venv", "venv", "node_modules", ".eggs"}
        
        for py_file in root.rglob("*.py"):
            # Check depth
            rel = py_file.relative_to(root)
            if len(rel.parts) > max_depth + 1:
                continue
            
            # Skip ignored directories
            if any(part in ignore for part in rel.parts):
                continue
            
            try:
                mod_info = self._parse_module(py_file, root)
                if mod_info:
                    modules[mod_info.name] = mod_info
                    self._module_cache[mod_info.name] = mod_info
            except Exception as e:
                logger.debug(f"Could not parse {py_file}: {e}")
        
        return modules
    
    def _parse_module(self, file_path: Path, root: Path) -> Optional[ModuleInfo]:
        """Parse a single Python module"""
        try:
            source = file_path.read_text(encoding='utf-8', errors='ignore')
            tree = ast.parse(source)
        except SyntaxError:
            return None
        
        rel = file_path.relative_to(root)
        name = str(rel).replace(os.sep, '.').replace('.py', '')
        if name.endswith('.__init__'):
            name = name[:-9]
        
        mod = ModuleInfo(
            name=name,
            path=str(file_path),
            line_count=len(source.splitlines()),
        )
        
        # Extract docstring
        if (tree.body and isinstance(tree.body[0], ast.Expr) and 
            isinstance(tree.body[0].value, ast.Constant) and 
            isinstance(tree.body[0].value.value, str)):
            mod.docstring = tree.body[0].value.value.strip().split('\n')[0]
        
        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                mod.classes.append(node.name)
            elif isinstance(node, ast.FunctionDef) and isinstance(node, ast.FunctionDef):
                # Only top-level functions
                if any(node is child for child in ast.iter_child_nodes(tree)):
                    mod.functions.append(node.name)
            elif isinstance(node, ast.Import):
                for alias in node.names:
                    mod.imports.append(alias.name)
            elif isinstance(node, ast.ImportFrom):
                if node.module:
                    mod.imports.append(node.module)
        
        return mod
    
    def _extract_classes(self, files: List[str]) -> Dict[str, ClassInfo]:
        """Extract class information from files"""
        classes = {}
        
        for file_path in files:
            try:
                source = Path(file_path).read_text(encoding='utf-8', errors='ignore')
                tree = ast.parse(source)
            except Exception:
                continue
            
            module_name = Path(file_path).stem
            
            for node in ast.walk(tree):
                if isinstance(node, ast.ClassDef):
                    cls = ClassInfo(
                        name=node.name,
                        module=module_name,
                        bases=[self._get_name(b) for b in node.bases],
                        docstring=ast.get_docstring(node) or "",
                    )
                    
                    for item in node.body:
                        if isinstance(item, ast.FunctionDef):
                            cls.methods.append(item.name)
                        elif isinstance(item, ast.Assign):
                            for target in item.targets:
                                if isinstance(target, ast.Name):
                                    cls.attributes.append(target.id)
                    
                    classes[node.name] = cls
                    self._class_cache[node.name] = cls
        
        return classes
    
    def _extract_calls(self, tree: ast.AST, function_name: str) -> List[Tuple[str, str, str]]:
        """Extract function calls within a specific function"""
        calls = []
        
        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                if node.name == function_name:
                    for child in ast.walk(node):
                        if isinstance(child, ast.Call):
                            callee = self._get_call_name(child)
                            if callee:
                                calls.append((function_name, callee, ""))
        
        return calls
    
    @staticmethod
    def _get_name(node) -> str:
        """Get name from an AST node"""
        if isinstance(node, ast.Name):
            return node.id
        elif isinstance(node, ast.Attribute):
            return f"{LivingDocumentation._get_name(node.value)}.{node.attr}"
        return "unknown"
    
    @staticmethod
    def _get_call_name(node: ast.Call) -> Optional[str]:
        """Get the function name from a Call node"""
        if isinstance(node.func, ast.Name):
            return node.func.id
        elif isinstance(node.func, ast.Attribute):
            return node.func.attr
        return None
    
    @staticmethod
    def _sanitize_id(name: str) -> str:
        """Sanitize a name for use as a Mermaid node ID"""
        return re.sub(r'[^a-zA-Z0-9_]', '_', name)
    
    @staticmethod
    def _write_diagram(path: str, title: str, diagram: str):
        """Write a diagram to a markdown file"""
        content = f"""# {title}

*Auto-generated by OmniClaw LivingDocumentation — {time.strftime('%Y-%m-%d %H:%M:%S')}*

```mermaid
{diagram}
```
"""
        with open(path, 'w', encoding='utf-8') as f:
            f.write(content)
    
    def get_stats(self) -> Dict[str, Any]:
        """Get documentation statistics"""
        return {
            "cached_modules": len(self._module_cache),
            "cached_classes": len(self._class_cache),
            "last_scan": time.strftime('%Y-%m-%d %H:%M:%S', 
                                       time.localtime(self._last_scan)) if self._last_scan else "never",
        }
