"""Ponytail MCP Server — help agents simplify code."""
import ast
import os
import re
import sys
from pathlib import Path
from typing import Any

from mcp.server import Server, stdio_server
from mcp.types import TextContent, Tool

# --- helpers ---

PYTHON_KEYWORDS = frozenset(
    "False None True and as assert async await break class continue "
    "def del elif else except finally for from global if import in "
    "is lambda nonlocal not or pass raise return try while with yield".split()
)


def walk_python_files(root: str) -> list[Path]:
    """Recursively find all .py files under root."""
    found = []
    for p in Path(root).rglob("*.py"):
        try:
            rel = p.relative_to(root)
        except ValueError:
            rel = p
        found.append(rel)
    return sorted(found)


def read_lines(path: str) -> list[str]:
    try:
        with open(path) as f:
            return f.readlines()
    except OSError:
        return []


# --- tool 1: find reinvented stdlib ---

# patterns: function names / class names that reimplement stdlib
REINVENTION_PATTERNS = {
    "flatten": re.compile(r"def\s+(flatten|flatten_list)\s*\("),
    "average": re.compile(r"def\s+(average|avg|mean)\s*\("),
    "url_parse": re.compile(
        r"(def\s+(parse_url|url_parse|parse_uri|url_parser)\s*\(|"
        r"class\s+(URLParser|URIParser|UrlParser))"
    ),
    "deep_merge": re.compile(r"def\s+(deep_merge|merge_dicts)\s*\("),
    "chunk": re.compile(r"def\s+(chunk|chunked|chunks)\s*\("),
    "debounce": re.compile(r"def\s+(debounce|throttle)\s*\("),
    "slugify": re.compile(r"def\s+(slugify|to_slug)\s*\("),
    "unique": re.compile(r"def\s+(unique|deduplicate|distinct)\s*\("),
    "flatten_json": re.compile(r"def\s+(flatten_json|flatten_dict)\s*\("),
    "retry": re.compile(r"def\s+(retry|retry_with|with_retry)\s*\("),
}


def tool_find_reinvented_stdlib(_params: dict[str, Any]) -> list[dict[str, Any]]:
    """Find reimplemented Python standard library functions."""
    cwd = os.getcwd()
    results: list[dict[str, Any]] = []
    for pyfile in walk_python_files(cwd):
        lines = read_lines(os.path.join(cwd, str(pyfile)))
        for i, line in enumerate(lines, 1):
            stripped = line.strip()
            if stripped.startswith("#") or stripped.startswith("'''") or stripped.startswith('"""'):
                continue
            for name, pattern in REINVENTION_PATTERNS.items():
                if pattern.search(line):
                    results.append(
                        {"file": str(pyfile), "line": i, "code": stripped}
                    )
                    break
    return results


# --- tool 2: find dead code ---

def _collect_imports(root: str, pyfiles: list[Path]) -> dict[str, str]:
    """Map module name -> relative path."""
    imports: dict[str, str] = {}
    for pyfile in pyfiles:
        src = "\n".join(read_lines(os.path.join(root, str(pyfile))))
        try:
            tree = ast.parse(src, filename=str(pyfile))
        except SyntaxError:
            continue
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    mod = alias.asname or alias.name
                    imports[mod] = str(pyfile)
            elif isinstance(node, ast.ImportFrom):
                mod = node.module or ""
                if node.level > 0:
                    parts = str(pyfile).split("/")[: -node.level]
                    mod = "/".join(parts) + ("/" + mod if mod else "")
                imports[mod] = str(pyfile)
    return imports


def tool_find_dead_code(_params: dict[str, Any]) -> list[dict[str, Any]]:
    """Find files never imported and symbols never referenced."""
    cwd = os.getcwd()
    pyfiles = walk_python_files(cwd)
    if not pyfiles:
        return []

    # Build set of all imported module names
    imports = _collect_imports(cwd, pyfiles)

    results: list[dict[str, Any]] = []

    # Dead files: .py files not imported anywhere (excluding __init__)
    for pyfile in pyfiles:
        name = pyfile.stem
        if name == "__init__":
            continue
        stem = pyfile.stem
        if stem not in imports:
            # Check if any file imports this module
            found_import = False
            for pyfile2 in pyfiles:
                if pyfile2 == pyfile:
                    continue
                src = "\n".join(read_lines(os.path.join(cwd, str(pyfile2))))
                try:
                    tree = ast.parse(src, filename=str(pyfile2))
                except SyntaxError:
                    continue
                for node in ast.walk(tree):
                    if isinstance(node, ast.Import):
                        for alias in node.names:
                            if alias.name == stem or alias.name.startswith(stem + "."):
                                found_import = True
                                break
                    elif isinstance(node, ast.ImportFrom):
                        if node.module and (node.module == stem or node.module.startswith(stem + ".")):
                            found_import = True
                            break
            if not found_import:
                results.append(
                    {
                        "file": str(pyfile),
                        "line": 1,
                        "symbol": stem,
                        "reason": "module never imported",
                    }
                )

    # Dead symbols: top-level functions/classes not referenced elsewhere
    for pyfile in pyfiles:
        src = "\n".join(read_lines(os.path.join(cwd, str(pyfile))))
        try:
            tree = ast.parse(src, filename=str(pyfile))
        except SyntaxError:
            continue

        # Collect defined top-level names
        defined: list[tuple[str, int]] = []
        for node in ast.iter_child_nodes(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                defined.append((node.name, node.lineno))
            elif isinstance(node, ast.ClassDef):
                defined.append((node.name, node.lineno))

        # Collect all Name references across the whole project
        all_refs: set[str] = set()
        for other in pyfiles:
            other_src = "\n".join(read_lines(os.path.join(cwd, str(other))))
            try:
                other_tree = ast.parse(other_src, filename=str(other))
            except SyntaxError:
                continue
            for n in ast.walk(other_tree):
                if isinstance(n, ast.Name):
                    all_refs.add(n.id)
                elif isinstance(n, ast.Attribute):
                    if isinstance(n.value, ast.Name):
                        all_refs.add(n.value.id)

        # Also check self-references within same file (imports count as refs)
        for imp_name in imports:
            base = imp_name.split(".")[-1]
            all_refs.add(base)

        for sym_name, lineno in defined:
            if sym_name.startswith("_"):
                continue
            if sym_name in PYTHON_KEYWORDS:
                continue
            # Check if referenced outside its own file
            refs_outside = False
            for other in pyfiles:
                if other == pyfile:
                    continue
                other_src = "\n".join(read_lines(os.path.join(cwd, str(other))))
                try:
                    other_tree = ast.parse(other_src, filename=str(other))
                except SyntaxError:
                    continue
                for n in ast.walk(other_tree):
                    if isinstance(n, ast.Name) and n.id == sym_name:
                        refs_outside = True
                        break
                    if isinstance(n, ast.Attribute) and isinstance(n.value, ast.Name) and n.value.id == sym_name:
                        refs_outside = True
                        break
                if refs_outside:
                    break

            if not refs_outside:
                results.append(
                    {
                        "file": str(pyfile),
                        "line": lineno,
                        "symbol": sym_name,
                        "reason": "symbol never referenced",
                    }
                )

    return results


# --- tool 3: find single-impl interfaces ---

def tool_find_single_impl_interface(_params: dict[str, Any]) -> list[dict[str, Any]]:
    """Find abstract classes with exactly one concrete implementation."""
    cwd = os.getcwd()
    pyfiles = walk_python_files(cwd)
    if not pyfiles:
        return []

    # Collect all abstract base classes
    abstract_classes: dict[str, tuple[str, int]] = {}
    concrete_implementations: dict[str, list[tuple[str, int]]] = {}

    for pyfile in pyfiles:
        src = "\n".join(read_lines(os.path.join(cwd, str(pyfile))))
        try:
            tree = ast.parse(src, filename=str(pyfile))
        except SyntaxError:
            continue

        for node in ast.iter_child_nodes(tree):
            if not isinstance(node, ast.ClassDef):
                continue

            # Check if it's an abstract class
            is_abstract = False
            for base in node.bases:
                base_name = ""
                if isinstance(base, ast.Name):
                    base_name = base.id
                elif isinstance(base, ast.Attribute):
                    base_name = base.attr

                if base_name in ("ABC", "ABCMeta"):
                    is_abstract = True
                if isinstance(base, ast.Attribute):
                    full = f"{ast.dump(base)}"
                    if "ABCMeta" in full or "ABC" in full:
                        is_abstract = True

            # Check for abstract methods (decorated with @abstractmethod)
            has_abstract_methods = False
            for child in ast.iter_child_nodes(node):
                if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    for dec in child.decorator_list:
                        dec_name = ""
                        if isinstance(dec, ast.Name):
                            dec_name = dec.id
                        elif isinstance(dec, ast.Attribute):
                            dec_name = dec.attr
                        if "abstractmethod" in dec_name.lower():
                            has_abstract_methods = True
                            break

            class_name = node.name
            if is_abstract or has_abstract_methods:
                abstract_classes[class_name] = (str(pyfile), node.lineno)
                concrete_implementations[class_name] = []

            # Check if this class inherits from a known abstract class
            for base in node.bases:
                base_name = ""
                if isinstance(base, ast.Name):
                    base_name = base.id
                elif isinstance(base, ast.Attribute):
                    base_name = base.attr
                if base_name in abstract_classes:
                    concrete_implementations[base_name].append(
                        (str(pyfile), node.lineno)
                    )

    # Also check for ABCMeta metaclass usage
    for pyfile in pyfiles:
        src = "\n".join(read_lines(os.path.join(cwd, str(pyfile))))
        try:
            tree = ast.parse(src, filename=str(pyfile))
        except SyntaxError:
            continue
        for node in ast.iter_child_nodes(tree):
            if isinstance(node, ast.ClassDef):
                for kw in node.keywords:
                    if kw.arg == "metaclass":
                        if isinstance(kw.value, ast.Name) and kw.value.id == "ABCMeta":
                            cn = node.name
                            if cn not in abstract_classes:
                                abstract_classes[cn] = (str(pyfile), node.lineno)
                                concrete_implementations[cn] = []

    results: list[dict[str, Any]] = []
    for cls_name, (file_path, line_no) in abstract_classes.items():
        impls = concrete_implementations.get(cls_name, [])
        if len(impls) == 1:
            impl_file, impl_line = impls[0]
            results.append(
                {
                    "abstract": cls_name,
                    "implementation": impls[0][0].rsplit("/", 1)[-1].replace(".py", ""),
                    "file": impl_file,
                    "line": impl_line,
                }
            )
    return results


# --- tool 4: find ponytail comments ---

def tool_find_ponytail_comments(_params: dict[str, Any]) -> list[dict[str, Any]]:
    """Extract all //ponytail: or #ponytail: comments."""
    cwd = os.getcwd()
    results: list[dict[str, Any]] = []
    ponytail_re = re.compile(r"(?://|#)ponytail:\s*(.*)")

    for pyfile in walk_python_files(cwd):
        lines = read_lines(os.path.join(cwd, str(pyfile)))
        for i, line in enumerate(lines, 1):
            # Skip string literals (simple check)
            stripped = line.strip()
            if stripped.startswith("#!") or stripped.startswith('"""') or stripped.startswith("'''"):
                continue
            match = ponytail_re.search(line)
            if match:
                results.append(
                    {
                        "file": str(pyfile),
                        "line": i,
                        "comment": match.group(1).strip(),
                    }
                )
    return results


# --- tool 5: analyze bloat ---

def _nesting_depth(node: ast.AST) -> int:
    """Compute max nesting depth of control structures."""
    max_depth = 0
    for child in ast.iter_child_nodes(node):
        if isinstance(child, (ast.For, ast.AsyncFor, ast.While, ast.If, ast.With,
                              ast.AsyncWith, ast.Try, ast.ExceptHandler)):
            child_depth = _nesting_depth(child) + 1
            max_depth = max(max_depth, child_depth)
        else:
            child_depth = _nesting_depth(child)
            max_depth = max(max_depth, child_depth)
    return max_depth


def _func_lengths(tree: ast.Module) -> list[int]:
    """Get lengths of all top-level and nested functions."""
    lengths: list[int] = []

    def walk(node: ast.AST):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            end = max(
                (getattr(n, "end_lineno", n.lineno) for n in ast.walk(node)),
                default=node.lineno,
            )
            lengths.append(end - node.lineno + 1)
        for child in ast.iter_child_nodes(node):
            walk(child)

    walk(tree)
    return lengths


def tool_analyze_bloat(params: dict[str, Any]) -> dict[str, Any]:
    """Analyze a file for bloat: lines, function lengths, nesting, imports."""
    filepath = params.get("path", "")
    if not filepath:
        return {"error": "path required"}

    if not os.path.exists(filepath):
        return {"error": f"file not found: {filepath}"}

    lines = read_lines(filepath)
    total_lines = len(lines)
    if total_lines == 0:
        return {
            "file": filepath,
            "total_lines": 0,
            "avg_function_length": 0,
            "max_nesting_depth": 0,
            "import_count": 0,
        }

    src = "\n".join(lines)
    try:
        tree = ast.parse(src, filename=filepath)
    except SyntaxError:
        return {
            "file": filepath,
            "total_lines": total_lines,
            "avg_function_length": 0,
            "max_nesting_depth": 0,
            "import_count": 0,
            "parse_error": True,
        }

    func_lengths = _func_lengths(tree)
    avg_func_len = (
        sum(func_lengths) / len(func_lengths) if func_lengths else 0
    )

    max_nesting = _nesting_depth(tree)

    import_count = 0
    for node in ast.walk(tree):
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            import_count += 1

    return {
        "file": filepath,
        "total_lines": total_lines,
        "avg_function_length": round(avg_func_len, 1),
        "max_nesting_depth": max_nesting,
        "import_count": import_count,
    }


# --- server setup ---

async def run_server() -> None:
    """Run the ponytail MCP server."""
    server = Server("ponytail")

    @server.list_tools()
    async def list_tools() -> list[Tool]:
        return [
            Tool(
                name="ponytail_find_reinvented_stdlib",
                description="Find reimplemented Python standard library functions (flatten, average, URL parser, etc.)",
                inputSchema={
                    "type": "object",
                    "properties": {},
                },
            ),
            Tool(
                name="ponytail_find_dead_code",
                description="Find files never imported and symbols never referenced",
                inputSchema={
                    "type": "object",
                    "properties": {},
                },
            ),
            Tool(
                name="ponytail_find_single_impl_interface",
                description="Find abstract classes with exactly one concrete implementation",
                inputSchema={
                    "type": "object",
                    "properties": {},
                },
            ),
            Tool(
                name="ponytail_find_ponytail_comments",
                description="Extract all //ponytail: or #ponytail: comments from the workspace",
                inputSchema={
                    "type": "object",
                    "properties": {},
                },
            ),
            Tool(
                name="ponytail_analyze_bloat",
                description="Analyze a file for bloat metrics: total lines, avg function length, max nesting depth, import count",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Path to the Python file to analyze",
                        }
                    },
                    "required": ["path"],
                },
            ),
        ]

    @server.call_tool()
    async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
        handlers = {
            "ponytail_find_reinvented_stdlib": tool_find_reinvented_stdlib,
            "ponytail_find_dead_code": tool_find_dead_code,
            "ponytail_find_single_impl_interface": tool_find_single_impl_interface,
            "ponytail_find_ponytail_comments": tool_find_ponytail_comments,
            "ponytail_analyze_bloat": tool_analyze_bloat,
        }
        handler = handlers.get(name)
        if not handler:
            return [TextContent(type="text", text=f"Unknown tool: {name}")]
        result = handler(arguments or {})
        return [TextContent(type="text", text=__import__("json").dumps(result, separators=(",", ":")))]

    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream)


if __name__ == "__main__":
    import asyncio
    asyncio.run(run_server())