"""
GitMem Coding - AST Skeleton Generator

Generates compact, token-efficient code skeletons from source files
using Python's built-in `ast` module. Inspired by Aider's repomap technique.

The skeleton preserves:
  - Module docstring
  - Import statements
  - Top-level constants/assignments
  - Class definitions (name, bases, docstring, attributes, method signatures)
  - Function definitions (name, full signature, docstring first line)
  - Decorators

The skeleton omits:
  - Function/method bodies (replaced with ⋮...)
  - Inline comments (except docstrings)
  - Blank lines, formatting details

Typical compression: 70-85% token reduction vs full source.
"""

import ast
import os
import re
from typing import List, Optional, Tuple


# ─── Helpers ────────────────────────────────────────────────────────────────

ELLIPSIS_MARKER = "    ⋮..."

# Map file extensions → language identifier
EXTENSION_TO_LANGUAGE = {
    ".py": "python",
    ".js": "javascript",
    ".ts": "typescript",
    ".tsx": "typescript",
    ".jsx": "javascript",
    ".java": "java",
    ".go": "go",
    ".rs": "rust",
    ".rb": "ruby",
    ".php": "php",
    ".c": "c",
    ".cpp": "cpp",
    ".h": "c",
    ".hpp": "cpp",
    ".cs": "csharp",
    ".swift": "swift",
    ".kt": "kotlin",
    ".html": "html",
    ".css": "css",
    ".sql": "sql",
    ".sh": "shell",
    ".bash": "shell",
    ".yaml": "yaml",
    ".yml": "yaml",
    ".json": "json",
    ".toml": "toml",
    ".md": "markdown",
    ".xml": "xml",
}


def detect_language(file_path: str) -> str:
    """Auto-detect language from file extension."""
    ext = os.path.splitext(file_path)[1].lower()
    return EXTENSION_TO_LANGUAGE.get(ext, "other")


def _truncate_docstring(docstring: str, max_chars: int = 2000) -> str:
    """
    Format a docstring for the skeleton. 
    Preserves the full content (up to max_chars) to capture semantic meaning.
    """
    if not docstring:
        return ""
    
    doc = docstring.strip()
    if len(doc) <= max_chars:
        return doc
    
    return doc[:max_chars] + "..."


# ─── Python AST Skeleton ───────────────────────────────────────────────────

class PythonSkeletonGenerator:
    """
    Generates compact skeletons from Python source using the `ast` module.
    
    Extracts a structural outline: imports, class/function signatures,
    docstrings, decorators, and class-level attributes.
    """

    def generate(self, source: str, file_path: str = "") -> str:
        """
        Generate a compact skeleton from Python source code.
        
        Args:
            source: Full Python source code
            file_path: Optional file path for the header line
            
        Returns:
            Compact skeleton string
        """
        try:
            tree = ast.parse(source)
        except SyntaxError:
            # Fallback to generic if Python parse fails
            return GenericSkeletonGenerator().generate(source, file_path)

        lines = source.splitlines()
        line_count = len(lines)
        file_name = os.path.basename(file_path) if file_path else "unknown"

        parts: List[str] = []

        # ── Header ──
        parts.append(f"# {file_name} | {line_count} lines | python")

        # ── Module docstring ──
        mod_doc = ast.get_docstring(tree)
        if mod_doc:
            parts.append(f'"""\n{mod_doc}\n"""')

        # ── Gather top-level nodes ──
        imports: List[str] = []
        constants: List[str] = []
        definitions: List[str] = []

        for node in tree.body:
            # Skip the module docstring expression (already handled)
            if (isinstance(node, ast.Expr) and
                    isinstance(node.value, (ast.Constant,))):
                continue

            if isinstance(node, (ast.Import, ast.ImportFrom)):
                imports.append(self._format_import(node))

            elif isinstance(node, ast.Assign):
                const_line = self._format_assignment(node, lines)
                if const_line:
                    constants.append(const_line)

            elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                definitions.append(self._format_function(node, lines, indent=0))

            elif isinstance(node, ast.ClassDef):
                definitions.append(self._format_class(node, lines))

        # ── Assemble ──
        if imports:
            parts.append("")
            parts.extend(imports)

        if constants:
            parts.append("")
            parts.extend(constants)

        if definitions:
            for defn in definitions:
                parts.append("")
                parts.append(defn)

        return "\n".join(parts)

    # ── Formatters ──────────────────────────────────────────────────────

    def _format_import(self, node: ast.AST) -> str:
        """Format an import statement from AST."""
        if isinstance(node, ast.Import):
            names = ", ".join(
                f"{a.name} as {a.asname}" if a.asname else a.name
                for a in node.names
            )
            return f"import {names}"
        elif isinstance(node, ast.ImportFrom):
            module = node.module or ""
            dots = "." * (node.level or 0)
            names = ", ".join(
                f"{a.name} as {a.asname}" if a.asname else a.name
                for a in node.names
            )
            return f"from {dots}{module} import {names}"
        return ""

    def _format_assignment(self, node: ast.Assign, lines: List[str]) -> Optional[str]:
        """Format a top-level assignment (constant)."""
        # Only include simple name = value assignments (not unpacking)
        if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
            name = node.targets[0].id
            # Skip private/dunder unless they look important
            if name.startswith("__") and name.endswith("__"):
                # Include __version__, __all__, etc.
                if name not in ("__version__", "__all__", "__author__"):
                    return None
            # Get the source line for value
            try:
                line = lines[node.lineno - 1].strip()
                # Truncate very long assignments
                if len(line) > 120:
                    line = line[:117] + "..."
                return line
            except IndexError:
                return f"{name} = ..."
        return None

    def _format_function(self, node, lines: List[str], indent: int = 0) -> str:
        """Format a function/method definition with its signature and docstring."""
        prefix = "    " * indent
        parts: List[str] = []

        # Decorators
        for dec in node.decorator_list:
            dec_text = self._get_decorator_text(dec, lines)
            parts.append(f"{prefix}@{dec_text}")

        # Signature
        is_async = isinstance(node, ast.AsyncFunctionDef)
        sig = self._build_signature(node)
        keyword = "async def" if is_async else "def"
        parts.append(f"{prefix}{keyword} {node.name}{sig}:")

        # Docstring
        doc = ast.get_docstring(node)
        if doc:
            parts.append(f'{prefix}    """\n{prefix}    {doc}\n{prefix}    """')

        # Body analysis
        for stmt in node.body:
            # Skip docstrings (already handled)
            if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Constant):
                 # Check if constant is a string (docstring)
                 if isinstance(stmt.value.value, str):
                     continue
            
            # Check for input extraction
            snippet = self._format_input_extraction(stmt, lines)
            if snippet:
                parts.append(f"{prefix}    {snippet}")

        # Body placeholder
        parts.append(f"{prefix}{ELLIPSIS_MARKER}")

        return "\n".join(parts)

    def _format_input_extraction(self, node: ast.AST, lines: List[str]) -> Optional[str]:
        """
        Check if a statement is an input extraction and return its source code.
        Looks for assignments involving 'request.', 'args', 'kwargs', etc.
        """
        if not isinstance(node, (ast.Assign, ast.AnnAssign)):
            return None
            
        # Get source text
        try:
            # Handle potential missing end_lineno in older Python
            end_line = getattr(node, 'end_lineno', node.lineno)
            # Extract lines (0-indexed)
            source_lines = lines[node.lineno-1:end_line]
            source_text = "\n".join(source_lines).strip()
            
            # Check for keywords
            # We look for "request." (Flask/Django), "kwargs.get", "sys.argv"
            # simple string matching is usually effective enough and safer than deep AST inspection
            keywords = ["request.", "kwargs.get", "sys.argv", "argparse", "optparse"]
            if any(k in source_text for k in keywords):
                return source_text
                
        except (IndexError, AttributeError):
            pass
            
        return None

    def _format_class(self, node: ast.ClassDef, lines: List[str]) -> str:
        """Format a class definition with attributes and method signatures."""
        parts: List[str] = []

        # Decorators
        for dec in node.decorator_list:
            dec_text = self._get_decorator_text(dec, lines)
            parts.append(f"@{dec_text}")

        # Class header
        bases = ", ".join(self._get_node_name(b) for b in node.bases)
        if bases:
            parts.append(f"class {node.name}({bases}):")
        else:
            parts.append(f"class {node.name}:")

        # Docstring
        doc = ast.get_docstring(node)
        if doc:
            parts.append(f'    """\n    {doc}\n    """')

        # Class-level attributes (assignments in class body)
        for item in node.body:
            # Skip the docstring expression
            if (isinstance(item, ast.Expr) and
                    isinstance(item.value, (ast.Constant,))):
                continue

            if isinstance(item, ast.Assign):
                attr = self._format_class_attr(item, lines)
                if attr:
                    parts.append(f"    {attr}")

            elif isinstance(item, ast.AnnAssign):
                attr = self._format_annotated_attr(item, lines)
                if attr:
                    parts.append(f"    {attr}")

            elif isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
                parts.append("")
                parts.append(self._format_function(item, lines, indent=1))

            elif isinstance(item, ast.ClassDef):
                # Nested class — just show name
                parts.append(f"    class {item.name}: ...")

        return "\n".join(parts)

    def _format_class_attr(self, node: ast.Assign, lines: List[str]) -> Optional[str]:
        """Format a class-level attribute assignment."""
        if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
            try:
                line = lines[node.lineno - 1].strip()
                if len(line) > 100:
                    line = line[:97] + "..."
                return line
            except IndexError:
                return None
        return None

    def _format_annotated_attr(self, node: ast.AnnAssign, lines: List[str]) -> Optional[str]:
        """Format an annotated class attribute."""
        try:
            line = lines[node.lineno - 1].strip()
            if len(line) > 100:
                line = line[:97] + "..."
            return line
        except IndexError:
            return None

    def _build_signature(self, node) -> str:
        """Build a function signature string from AST node."""
        args = node.args
        parts: List[str] = []

        # Positional-only args
        posonlyargs = getattr(args, 'posonlyargs', [])
        
        # Regular args
        all_args = posonlyargs + args.args
        num_defaults = len(args.defaults)
        num_args = len(all_args)

        for i, arg in enumerate(all_args):
            arg_str = arg.arg
            # Type annotation
            if arg.annotation:
                ann = self._get_node_name(arg.annotation)
                arg_str += f": {ann}"
            # Default value
            default_idx = i - (num_args - num_defaults)
            if default_idx >= 0:
                default = args.defaults[default_idx]
                default_str = self._get_value_repr(default)
                arg_str += f" = {default_str}"
            parts.append(arg_str)
            # Insert / after positional-only args
            if i == len(posonlyargs) - 1 and posonlyargs:
                parts.append("/")

        # *args
        if args.vararg:
            vararg_str = f"*{args.vararg.arg}"
            if args.vararg.annotation:
                vararg_str += f": {self._get_node_name(args.vararg.annotation)}"
            parts.append(vararg_str)
        elif args.kwonlyargs:
            parts.append("*")

        # Keyword-only args
        for i, arg in enumerate(args.kwonlyargs):
            arg_str = arg.arg
            if arg.annotation:
                ann = self._get_node_name(arg.annotation)
                arg_str += f": {ann}"
            if i < len(args.kw_defaults) and args.kw_defaults[i] is not None:
                default_str = self._get_value_repr(args.kw_defaults[i])
                arg_str += f" = {default_str}"
            parts.append(arg_str)

        # **kwargs
        if args.kwarg:
            kwarg_str = f"**{args.kwarg.arg}"
            if args.kwarg.annotation:
                kwarg_str += f": {self._get_node_name(args.kwarg.annotation)}"
            parts.append(kwarg_str)

        sig = f"({', '.join(parts)})"

        # Return annotation
        if node.returns:
            ret = self._get_node_name(node.returns)
            sig += f" -> {ret}"

        return sig

    def _get_node_name(self, node) -> str:
        """Get a readable name from an AST node (for types, bases, etc.)."""
        if node is None:
            return ""
        if isinstance(node, ast.Name):
            return node.id
        if isinstance(node, ast.Attribute):
            value = self._get_node_name(node.value)
            return f"{value}.{node.attr}" if value else node.attr
        if isinstance(node, ast.Subscript):
            value = self._get_node_name(node.value)
            slice_val = self._get_node_name(node.slice)
            return f"{value}[{slice_val}]"
        if isinstance(node, ast.Constant):
            return repr(node.value)
        if isinstance(node, ast.Tuple):
            elts = ", ".join(self._get_node_name(e) for e in node.elts)
            return elts
        if isinstance(node, ast.List):
            elts = ", ".join(self._get_node_name(e) for e in node.elts)
            return f"[{elts}]"
        if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):
            # Union type syntax: X | Y
            left = self._get_node_name(node.left)
            right = self._get_node_name(node.right)
            return f"{left} | {right}"
        # Fallback
        try:
            return ast.dump(node)
        except Exception:
            return "..."

    def _get_value_repr(self, node) -> str:
        """Get a compact repr of a default value node."""
        if isinstance(node, ast.Constant):
            r = repr(node.value)
            return r if len(r) <= 50 else r[:47] + "..."
        if isinstance(node, ast.Name):
            return node.id
        if isinstance(node, ast.Attribute):
            return self._get_node_name(node)
        if isinstance(node, (ast.List, ast.Tuple, ast.Set)):
            if not node.elts:
                return "[]" if isinstance(node, ast.List) else "()" if isinstance(node, ast.Tuple) else "set()"
            return "..."
        if isinstance(node, ast.Dict):
            return "{}" if not node.keys else "{...}"
        if isinstance(node, ast.Call):
            func = self._get_node_name(node.func)
            return f"{func}(...)"
        if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
            operand = self._get_value_repr(node.operand)
            return f"-{operand}"
        return "..."

    def _get_decorator_text(self, node, lines: List[str]) -> str:
        """Get decorator text from source lines."""
        try:
            line = lines[node.lineno - 1].strip()
            if line.startswith("@"):
                line = line[1:]
            # Truncate long decorators
            if len(line) > 80:
                line = line[:77] + "..."
            return line
        except IndexError:
            return "..."


# ─── Generic/Fallback Skeleton ─────────────────────────────────────────────

class GenericSkeletonGenerator:
    """
    Regex-based skeleton generator for non-Python files.
    
    Extracts patterns that look like:
      - import/require/include statements
      - function/method definitions
      - class/struct/interface definitions
      - export statements
    """

    # Common patterns across languages
    IMPORT_PATTERNS = [
        re.compile(r'^(import\s+.+)$', re.MULTILINE),
        re.compile(r'^(from\s+.+\s+import\s+.+)$', re.MULTILINE),
        re.compile(r'^(const\s+.+\s*=\s*require\(.+\)).*$', re.MULTILINE),
        re.compile(r'^(#include\s+[<"].+[>"])$', re.MULTILINE),
        re.compile(r'^(using\s+.+;)$', re.MULTILINE),
        re.compile(r'^(require\s+.+)$', re.MULTILINE),
    ]

    FUNCTION_PATTERNS = [
        # JS/TS: function name(...) / async function / export function
        re.compile(r'^(\s*(?:export\s+)?(?:async\s+)?function\s+\w+\s*\([^)]*\)[^{]*)\{?\s*$', re.MULTILINE),
        # JS/TS arrow: const name = (...) =>
        re.compile(r'^(\s*(?:export\s+)?(?:const|let|var)\s+\w+\s*=\s*(?:async\s+)?\([^)]*\)\s*(?::\s*\w+)?\s*=>)', re.MULTILINE),
        # Go: func name(...)
        re.compile(r'^(func\s+(?:\([^)]+\)\s+)?\w+\s*\([^)]*\)[^{]*)\{?\s*$', re.MULTILINE),
        # Rust: fn name(...)
        re.compile(r'^(\s*(?:pub\s+)?(?:async\s+)?fn\s+\w+\s*(?:<[^>]*>)?\s*\([^)]*\)[^{]*)\{?\s*$', re.MULTILINE),
        # Java/C#: public void name(...)
        re.compile(r'^(\s*(?:public|private|protected|static|abstract|override|virtual|async|\s)*\s+\w+\s+\w+\s*\([^)]*\)[^{;]*)', re.MULTILINE),
        # Ruby: def name
        re.compile(r'^(\s*def\s+\w+(?:\([^)]*\))?)', re.MULTILINE),
    ]

    CLASS_PATTERNS = [
        re.compile(r'^(\s*(?:export\s+)?(?:abstract\s+)?class\s+\w+[^{]*)', re.MULTILINE),
        re.compile(r'^(\s*(?:pub\s+)?struct\s+\w+[^{]*)', re.MULTILINE),
        re.compile(r'^(\s*(?:pub\s+)?(?:async\s+)?trait\s+\w+[^{]*)', re.MULTILINE),
        re.compile(r'^(\s*interface\s+\w+[^{]*)', re.MULTILINE),
        re.compile(r'^(\s*(?:pub\s+)?enum\s+\w+[^{]*)', re.MULTILINE),
    ]

    def generate(self, source: str, file_path: str = "") -> str:
        """Generate a compact skeleton from non-Python source code."""
        file_name = os.path.basename(file_path) if file_path else "unknown"
        language = detect_language(file_path)
        line_count = source.count("\n") + 1

        parts: List[str] = [f"# {file_name} | {line_count} lines | {language}"]

        # Extract imports
        imports = set()
        for pattern in self.IMPORT_PATTERNS:
            for match in pattern.finditer(source):
                imp = match.group(1).strip()
                if len(imp) <= 120:
                    imports.add(imp)
        if imports:
            parts.append("")
            parts.extend(sorted(imports))

        # Extract classes/structs
        classes = []
        for pattern in self.CLASS_PATTERNS:
            for match in pattern.finditer(source):
                cls = match.group(1).strip()
                if len(cls) <= 120:
                    classes.append(cls)
        if classes:
            parts.append("")
            for cls in classes:
                parts.append(cls)
                parts.append(ELLIPSIS_MARKER)

        # Extract functions
        functions = []
        for pattern in self.FUNCTION_PATTERNS:
            for match in pattern.finditer(source):
                func = match.group(1).strip()
                if len(func) <= 120:
                    functions.append(func)
        if functions:
            parts.append("")
            for func in functions:
                parts.append(func)
                parts.append(ELLIPSIS_MARKER)

        # If nothing extracted, just show file header + first few non-empty lines
        if not imports and not classes and not functions:
            parts.append("")
            parts.append("# (no structural elements extracted)")
            non_empty = [l.rstrip() for l in source.splitlines() if l.strip()][:10]
            parts.extend(non_empty)
            if line_count > 10:
                parts.append(ELLIPSIS_MARKER)

        return "\n".join(parts)


# ─── Outline Generator ─────────────────────────────────────────────────────

class OutlineGenerator:
    """
    Ultra-compact outline: just names and line numbers.
    Even more compressed than a skeleton — ideal for navigation.
    
    Example output:
        # server.py | 1221 lines
        L89    def get_data_dir()
        L132   def _normalize_agent_id()
        L143   def create_memory()
        L363   class Coder
        L363     def __init__()
        L400     def create()
        L500     def run()
    """

    def generate(self, source: str, file_path: str = "") -> str:
        """Generate an ultra-compact outline."""
        file_name = os.path.basename(file_path) if file_path else "unknown"
        language = detect_language(file_path)
        
        if language == "python":
            return self._python_outline(source, file_name)
        return self._generic_outline(source, file_name, language)

    def _python_outline(self, source: str, file_name: str) -> str:
        """AST-based outline for Python files."""
        try:
            tree = ast.parse(source)
        except SyntaxError:
            return self._generic_outline(source, file_name, "python")

        line_count = source.count("\n") + 1
        parts: List[str] = [f"# {file_name} | {line_count} lines"]

        for node in tree.body:
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                prefix = "async " if isinstance(node, ast.AsyncFunctionDef) else ""
                parts.append(f"L{node.lineno:<6} {prefix}def {node.name}()")
            elif isinstance(node, ast.ClassDef):
                parts.append(f"L{node.lineno:<6} class {node.name}")
                for item in node.body:
                    if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
                        prefix = "async " if isinstance(item, ast.AsyncFunctionDef) else ""
                        parts.append(f"L{item.lineno:<6}   {prefix}def {item.name}()")
            elif isinstance(node, ast.Assign):
                for target in node.targets:
                    if isinstance(target, ast.Name):
                        if target.id.isupper() or target.id.startswith("_"):
                            parts.append(f"L{node.lineno:<6} {target.id} = ...")

        return "\n".join(parts)

    def _generic_outline(self, source: str, file_name: str, language: str) -> str:
        """Regex-based outline for non-Python files."""
        line_count = source.count("\n") + 1
        parts: List[str] = [f"# {file_name} | {line_count} lines | {language}"]

        lines = source.splitlines()
        func_re = re.compile(r'(?:function|def|fn|func)\s+(\w+)')
        class_re = re.compile(r'(?:class|struct|trait|interface|enum)\s+(\w+)')

        for i, line in enumerate(lines, 1):
            stripped = line.strip()
            m = class_re.search(stripped)
            if m:
                parts.append(f"L{i:<6} class {m.group(1)}")
                continue
            m = func_re.search(stripped)
            if m:
                indent = len(line) - len(line.lstrip())
                prefix = "  " if indent > 0 else ""
                parts.append(f"L{i:<6} {prefix}def {m.group(1)}()")

        return "\n".join(parts)


# ─── Main API ──────────────────────────────────────────────────────────────

class ASTSkeletonGenerator:
    """
    Main entry point for skeleton generation.
    
    Dispatches to language-specific generators and provides
    token estimation for savings reporting.
    """

    def __init__(self):
        self._python_gen = PythonSkeletonGenerator()
        self._generic_gen = GenericSkeletonGenerator()
        self._outline_gen = OutlineGenerator()

    def generate_skeleton(
        self,
        source: str,
        language: str = "auto",
        file_path: str = ""
    ) -> str:
        """
        Generate a compact skeleton from source code.
        
        Args:
            source: Full source code content
            language: Programming language (or "auto" to detect from file_path)
            file_path: File path (used for auto-detection and header)
            
        Returns:
            Compact skeleton string
        """
        if language == "auto":
            language = detect_language(file_path)

        if language == "python":
            return self._python_gen.generate(source, file_path)
        return self._generic_gen.generate(source, file_path)

    def generate_outline(
        self,
        source: str,
        file_path: str = ""
    ) -> str:
        """
        Generate an ultra-compact outline (names + line numbers only).
        
        Args:
            source: Full source code content
            file_path: File path for header and language detection
            
        Returns:
            Ultra-compact outline string
        """
        return self._outline_gen.generate(source, file_path)

    @staticmethod
    def estimate_compression(original: str, skeleton: str) -> dict:
        """
        Calculate compression statistics.
        
        Returns:
            Dict with original_tokens, skeleton_tokens, compression_ratio, savings_pct
        """
        TOKENS_PER_CHAR = 0.25
        original_tokens = int(len(original) * TOKENS_PER_CHAR)
        skeleton_tokens = int(len(skeleton) * TOKENS_PER_CHAR)
        
        ratio = 1.0 - (skeleton_tokens / original_tokens) if original_tokens > 0 else 0.0
        
        return {
            "original_tokens": original_tokens,
            "skeleton_tokens": skeleton_tokens,
            "tokens_saved": original_tokens - skeleton_tokens,
            "compression_ratio": round(ratio, 4),
            "savings_pct": f"{ratio * 100:.1f}%"
        }
