"""albus init --scan -- Analyze source code and generate native albusOS skills.

Scans a Python codebase (any framework: MCP, FastAPI, Flask, plain scripts)
and generates:
  - skills/<namespace>/SKILL.md
  - skills/<namespace>/tools/<tool>.py  (run() wrappers)
  - albus.yaml agent configuration

Detection patterns (all via AST, no execution):
  1. @mcp.tool() decorated async functions
  2. @app.get/post/put/delete() decorated functions (FastAPI/Flask)
  3. Plain async def / def with type hints and docstrings
  4. Functions matching naming conventions (e.g. prefixed with verb)

The scanner is framework-agnostic — it works with any Python source.
"""

from __future__ import annotations

import ast
import logging
from pathlib import Path

from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


# =============================================================================
# Models
# =============================================================================


class DetectedParam(BaseModel):
    """A parameter extracted from a function signature."""

    name: str
    type_hint: str = "str"
    default: str | None = None
    description: str = ""
    required: bool = True


class DetectedTool(BaseModel):
    """A callable function detected as a potential tool."""

    name: str
    description: str
    params: list[DetectedParam] = Field(default_factory=list)
    source_file: str = ""
    source_module: str = ""
    decorator: str = ""  # e.g. "mcp.tool", "app.post", ""
    is_async: bool = True
    original_function: str = ""  # original function name in source
    return_type: str = "str"


class DetectedSkill(BaseModel):
    """A group of related tools forming a skill."""

    namespace: str
    description: str
    tools: list[DetectedTool] = Field(default_factory=list)
    source_files: list[str] = Field(default_factory=list)
    imports: list[str] = Field(default_factory=list)


# =============================================================================
# AST Scanner
# =============================================================================


def scan_directory(source_dir: Path) -> list[DetectedTool]:
    """Scan a directory for Python functions that look like tools.

    Walks all .py files, parses AST, and extracts functions that match
    tool patterns: decorated, async, documented, with typed parameters.

    When decorated tools are found (e.g. @mcp.tool(), @app.post()),
    non-decorated functions are filtered out — decorated functions are
    the definitive tool surface.

    Args:
        source_dir: Root directory to scan.

    Returns
    -------
        List of detected tools.
    """
    tools: list[DetectedTool] = []

    for py_file in sorted(source_dir.rglob("*.py")):
        if py_file.name.startswith("_") and py_file.name != "__init__.py":
            continue

        try:
            source = py_file.read_text(encoding="utf-8")
            tree = ast.parse(source, filename=str(py_file))
        except (SyntaxError, UnicodeDecodeError) as e:
            logger.debug("Skipping %s: %s", py_file, e)
            continue

        file_tools = _extract_tools_from_module(tree, py_file, source_dir)
        tools.extend(file_tools)

    # If any decorated tools exist, only keep decorated ones.
    # Decorated functions (@mcp.tool, @app.post, etc.) are the definitive
    # tool surface — non-decorated functions are internal helpers.
    has_decorated = any(t.decorator for t in tools)
    if has_decorated:
        tools = [t for t in tools if t.decorator]

    return tools


def _extract_tools_from_module(
    tree: ast.Module, py_file: Path, source_dir: Path
) -> list[DetectedTool]:
    """Extract tool-like functions from a parsed module."""
    tools: list[DetectedTool] = []
    rel_path = py_file.relative_to(source_dir)
    module_name = str(rel_path.with_suffix("")).replace("/", ".").replace("\\", ".")

    # Only look at top-level functions (skip class methods)
    for node in ast.iter_child_nodes(tree):
        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            continue

        # Skip private/dunder functions
        if node.name.startswith("_"):
            continue

        # Detect decorator pattern
        decorator = _detect_decorator(node)

        # Determine if this looks like a tool
        has_decorator = decorator != ""
        has_docstring = ast.get_docstring(node) is not None
        has_typed_params = _has_typed_params(node)
        is_async = isinstance(node, ast.AsyncFunctionDef)

        # Scoring: decorators are strong signals, docstrings + types are medium
        score = 0
        if has_decorator:
            score += 3
        if has_docstring:
            score += 2
        if has_typed_params:
            score += 1
        if is_async:
            score += 1

        # Need at least a decorator OR (docstring + typed params) to qualify
        if score < 3:
            continue

        # Extract tool info
        docstring = ast.get_docstring(node) or ""
        description = docstring.strip().split("\n")[0] if docstring else f"Tool: {node.name}"
        params = _extract_params(node, docstring)
        return_type = _extract_return_type(node)

        tool = DetectedTool(
            name=node.name,
            description=description,
            params=params,
            source_file=str(rel_path),
            source_module=module_name,
            decorator=decorator,
            is_async=is_async,
            original_function=node.name,
            return_type=return_type,
        )
        tools.append(tool)
        logger.debug("Detected tool: %s in %s", node.name, rel_path)

    return tools


def _detect_decorator(node: ast.FunctionDef | ast.AsyncFunctionDef) -> str:
    """Detect if a function has a tool/route decorator.

    Returns decorator string like "mcp.tool", "app.post", or "".
    """
    for dec in node.decorator_list:
        dec_str = _decorator_to_string(dec)
        if not dec_str:
            continue

        # MCP patterns: @mcp.tool(), @server.tool()
        if "tool" in dec_str.lower():
            return dec_str

        # FastAPI/Flask patterns: @app.get(), @router.post(), etc.
        for method in ("get", "post", "put", "delete", "patch"):
            if f".{method}" in dec_str.lower():
                return dec_str

    return ""


def _decorator_to_string(node: ast.expr) -> str:
    """Convert a decorator AST node to a readable string."""
    if isinstance(node, ast.Call):
        return _decorator_to_string(node.func)
    if isinstance(node, ast.Attribute):
        base = _decorator_to_string(node.value)
        return f"{base}.{node.attr}" if base else node.attr
    if isinstance(node, ast.Name):
        return node.id
    return ""


def _has_typed_params(node: ast.FunctionDef | ast.AsyncFunctionDef) -> bool:
    """Check if a function has type-annotated parameters."""
    for arg in node.args.args:
        if arg.arg in ("self", "cls", "context", "ctx"):
            continue
        if arg.annotation is not None:
            return True
    return False


def _extract_params(
    node: ast.FunctionDef | ast.AsyncFunctionDef, docstring: str
) -> list[DetectedParam]:
    """Extract parameters from function signature and docstring."""
    params: list[DetectedParam] = []
    doc_params = _parse_docstring_args(docstring)

    # Compute defaults alignment
    args = node.args
    num_args = len(args.args)
    num_defaults = len(args.defaults)
    first_default_idx = num_args - num_defaults

    for i, arg in enumerate(args.args):
        if arg.arg in ("self", "cls", "context", "ctx"):
            continue

        type_hint = _annotation_to_string(arg.annotation) if arg.annotation else "str"

        # Check if has default
        default_idx = i - first_default_idx
        default = None
        required = True
        if default_idx >= 0 and default_idx < len(args.defaults):
            default = _default_to_string(args.defaults[default_idx])
            required = False

        desc = doc_params.get(arg.arg, "")

        params.append(DetectedParam(
            name=arg.arg,
            type_hint=type_hint,
            default=default,
            description=desc,
            required=required,
        ))

    return params


def _extract_return_type(node: ast.FunctionDef | ast.AsyncFunctionDef) -> str:
    """Extract return type annotation as string."""
    if node.returns:
        return _annotation_to_string(node.returns)
    return "str"


def _annotation_to_string(node: ast.expr | None) -> str:
    """Convert a type annotation AST node to a string."""
    if node is None:
        return "str"
    if isinstance(node, ast.Constant):
        return str(node.value)
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        base = _annotation_to_string(node.value)
        return f"{base}.{node.attr}"
    if isinstance(node, ast.Subscript):
        base = _annotation_to_string(node.value)
        inner = _annotation_to_string(node.slice)
        return f"{base}[{inner}]"
    if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):
        left = _annotation_to_string(node.left)
        right = _annotation_to_string(node.right)
        return f"{left} | {right}"
    if isinstance(node, ast.Tuple):
        parts = [_annotation_to_string(e) for e in node.elts]
        return ", ".join(parts)
    return "str"


def _default_to_string(node: ast.expr) -> str:
    """Convert a default value AST node to a string."""
    if isinstance(node, ast.Constant):
        return repr(node.value)
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        base = _default_to_string(node.value)
        return f"{base}.{node.attr}"
    if isinstance(node, ast.List):
        return "[]"
    if isinstance(node, ast.Dict):
        return "{}"
    return "None"


def _parse_docstring_args(docstring: str) -> dict[str, str]:
    """Parse parameter descriptions from docstring (Google-style Args section)."""
    params: dict[str, str] = {}
    if not docstring:
        return params

    lines = docstring.split("\n")
    in_args = False

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

        if stripped.lower() in ("args:", "arguments:", "parameters:"):
            in_args = True
            continue

        if in_args:
            if stripped and not line.startswith((" ", "\t")) and ":" not in stripped[:30]:
                break
            if stripped.lower() in ("returns:", "raises:", "yields:", "examples:", "return:", "note:"):
                break

            if ":" in stripped:
                parts = stripped.split(":", 1)
                name = parts[0].strip()
                if "(" in name:
                    name = name.split("(")[0].strip()
                desc = parts[1].strip() if len(parts) > 1 else ""
                if name and not name.startswith("-") and not name.startswith("*"):
                    params[name] = desc

    return params


# =============================================================================
# Skill Grouping
# =============================================================================


def group_into_skills(tools: list[DetectedTool], namespace: str) -> list[DetectedSkill]:
    """Group detected tools into skills by source module.

    Each source file becomes a skill group. If there's only one source file
    (or a flat structure), everything goes into a single skill.

    Args:
        tools: List of detected tools.
        namespace: Skill namespace (usually inferred from source directory name).

    Returns
    -------
        List of detected skills.
    """
    if not tools:
        return []

    # Group by source file parent directory
    file_groups: dict[str, list[DetectedTool]] = {}
    for tool in tools:
        # Use the source file's parent module as grouping key
        source = Path(tool.source_file)
        group_key = source.parent.name if source.parent.name != "." else namespace
        file_groups.setdefault(group_key, []).append(tool)

    # If all tools are in one group or a flat structure, single skill
    if len(file_groups) <= 1:
        all_tools = [t for group in file_groups.values() for t in group]
        desc = _infer_skill_description(all_tools, namespace)
        source_files = list({t.source_file for t in all_tools})
        return [DetectedSkill(
            namespace=namespace,
            description=desc,
            tools=all_tools,
            source_files=source_files,
        )]

    # Multiple groups → multiple skills (or sub-skills under the namespace)
    skills: list[DetectedSkill] = []
    for group_name, group_tools in sorted(file_groups.items()):
        desc = _infer_skill_description(group_tools, group_name)
        source_files = list({t.source_file for t in group_tools})
        skills.append(DetectedSkill(
            namespace=namespace,
            description=desc,
            tools=group_tools,
            source_files=source_files,
        ))

    return skills


def _infer_skill_description(tools: list[DetectedTool], name: str) -> str:
    """Infer a skill description from its tools."""
    if len(tools) == 1:
        return tools[0].description
    tool_names = ", ".join(t.name for t in tools[:5])
    suffix = f" and {len(tools) - 5} more" if len(tools) > 5 else ""
    return f"{name.replace('_', ' ').title()} tools: {tool_names}{suffix}"


# =============================================================================
# Code Generation
# =============================================================================


def generate_skill_directory(
    skill: DetectedSkill,
    output_dir: Path,
    source_package: str,
) -> list[str]:
    """Generate a complete skill directory from detected tools.

    Creates:
      - skills/<namespace>/SKILL.md
      - skills/<namespace>/tools/<tool_name>.py

    Args:
        skill: The detected skill to generate.
        output_dir: Parent output directory (e.g. ./skills).
        source_package: Python import path to the source package
                        (e.g. "servicem8_mcp").

    Returns
    -------
        List of created file paths (relative).
    """
    skill_dir = output_dir / skill.namespace
    tools_dir = skill_dir / "tools"
    tools_dir.mkdir(parents=True, exist_ok=True)

    created: list[str] = []

    # Generate SKILL.md
    skill_md = _generate_skill_md(skill)
    (skill_dir / "SKILL.md").write_text(skill_md)
    created.append(str(skill_dir / "SKILL.md"))

    # Generate tool wrappers
    for tool in skill.tools:
        tool_file = tools_dir / f"{tool.name}.py"
        tool_code = _generate_tool_wrapper(tool, skill.namespace, source_package)
        tool_file.write_text(tool_code)
        created.append(str(tool_file))

    return created


def _generate_skill_md(skill: DetectedSkill) -> str:
    """Generate SKILL.md content for a skill."""
    lines = [
        f"# {skill.namespace.replace('_', ' ').title()}",
        "",
        skill.description,
        "",
        "## Tools",
        "",
    ]

    for tool in skill.tools:
        param_summary = ""
        required = [p for p in tool.params if p.required]
        if required:
            param_summary = " (" + ", ".join(p.name for p in required) + ")"
        lines.append(f"- **{skill.namespace}.{tool.name}**{param_summary} — {tool.description}")

    lines.extend([
        "",
        "## When to Use",
        "",
        f"Use this skill when working with {skill.namespace.replace('_', ' ')} operations.",
        "",
        f"*Auto-generated by `albus init --scan` from: {', '.join(skill.source_files)}*",
    ])

    return "\n".join(lines) + "\n"


def _generate_tool_wrapper(
    tool: DetectedTool, namespace: str, source_package: str
) -> str:
    """Generate a tool wrapper .py file that delegates to the original function."""
    # Build the docstring
    doc_lines = [f'"""{tool.description}"""']

    # Build imports
    lines = [
        doc_lines[0],
        "",
        "from __future__ import annotations",
        "",
        "from core.pathways.models.tool import ToolOutput",
        "",
    ]

    # Build run() signature
    sig_parts: list[str] = []
    for p in tool.params:
        type_str = _simplify_type(p.type_hint)
        if p.default is not None:
            sig_parts.append(f"{p.name}: {type_str} = {p.default}")
        else:
            sig_parts.append(f"{p.name}: {type_str}")
    sig_parts.append("context=None")

    sig = ", ".join(sig_parts)
    lines.append(f"async def run({sig}) -> ToolOutput:")

    # Build docstring for run()
    run_doc_lines = [f'    """{tool.description}']
    if tool.params:
        run_doc_lines.append("")
        run_doc_lines.append("    Args:")
        for p in tool.params:
            desc = p.description or f"{p.name} parameter"
            run_doc_lines.append(f"        {p.name}: {desc}")
    run_doc_lines.append('    """')
    lines.extend(run_doc_lines)

    # Build the body — delegate to original function
    # Import the original module and call the function
    lines.append(f"    from {source_package}.{_module_path(tool.source_module)} import {tool.original_function}")
    lines.append("")

    # Build call arguments
    call_args = ", ".join(
        f"{p.name}={p.name}" for p in tool.params
    )

    lines.append("    try:")
    if tool.is_async:
        lines.append(f"        result = await {tool.original_function}({call_args})")
    else:
        lines.append(f"        result = {tool.original_function}({call_args})")

    lines.append('        return ToolOutput(success=True, data={"result": str(result)})')
    lines.append("    except Exception as e:")
    lines.append("        return ToolOutput(success=False, error=str(e))")
    lines.append("")

    return "\n".join(lines)


def _module_path(source_module: str) -> str:
    """Convert source module path to import-friendly path.

    Strips the top-level package name since we import from source_package.
    e.g. "servicem8_mcp.tools.crud" → "tools.crud"
    """
    parts = source_module.split(".")
    if len(parts) > 1:
        return ".".join(parts[1:])
    return source_module


def _simplify_type(type_hint: str) -> str:
    """Simplify complex type hints for tool signatures."""
    # Map common patterns
    if type_hint in ("NoneType", "None"):
        return "str | None"
    return type_hint


# =============================================================================
# Config Generation
# =============================================================================


def generate_albus_yaml(
    skills: list[DetectedSkill],
    namespace: str,
) -> str:
    """Generate albus.yaml configuration for detected skills.

    Args:
        skills: List of detected skills.
        namespace: Root namespace (e.g. "servicem8").

    Returns
    -------
        YAML config string.
    """
    tool_names = []
    for skill in skills:
        for tool in skill.tools:
            tool_names.append(f"{skill.namespace}.{tool.name}")

    # Determine tool patterns
    tool_patterns = [f'"{namespace}.*"']

    # Count tools by rough category for agent suggestions
    total_tools = sum(len(s.tools) for s in skills)

    lines = [
        "# AlbusOS workspace configuration",
        "# Auto-generated by: albus init --scan",
        f"# Detected {total_tools} tools across {len(skills)} skill(s)",
        "#",
        "# Review and customize the agent configuration below.",
        "",
        "agents:",
        "  - id: host",
        f"    name: {namespace.replace('_', ' ').title()} Assistant",
        f'    tools: [{', '.join(tool_patterns)}, "memory.*"]',
        "    pathway: react",
        "    max_steps: 10",
        "    model: reasoning",
        "",
        "skills:",
        "  directories:",
        "    - ./skills",
        "",
        "# Model routing (optional — defaults built into albusOS)",
        "# models:",
        "#   routing:",
        "#     reasoning: openai/gpt-4o",
        "#     fast: openai/gpt-4o-mini",
        "",
        f"# Discovered tools ({total_tools}):",
    ]

    for skill in skills:
        for tool in skill.tools:
            lines.append(f"#   - {skill.namespace}.{tool.name}")

    return "\n".join(lines) + "\n"


# =============================================================================
# Main Entry Point
# =============================================================================


def run_scan(
    source_dir: Path,
    output_dir: Path | None = None,
    namespace: str | None = None,
) -> int:
    """Run the full scan pipeline.

    1. Scan source directory for tool-like functions
    2. Group into skills
    3. Generate skill directories + SKILL.md + tool wrappers
    4. Generate albus.yaml

    Args:
        source_dir: Directory to scan.
        output_dir: Workspace root (default: cwd).
        namespace: Skill namespace (default: inferred from source dir name).

    Returns
    -------
        Number of tools generated.
    """
    if output_dir is None:
        output_dir = Path.cwd()

    source_dir = source_dir.resolve()
    if not source_dir.exists():
        print(f"Error: Source directory does not exist: {source_dir}")
        return 0

    # Infer namespace from directory name
    if namespace is None:
        namespace = _infer_namespace(source_dir)

    # Infer source package for imports
    source_package = _infer_source_package(source_dir)

    print(f"Scanning: {source_dir}")
    print(f"Namespace: {namespace}")
    print(f"Source package: {source_package}")
    print()

    # 1. Scan
    tools = scan_directory(source_dir)
    if not tools:
        print("No tool-like functions detected.")
        print("  Looking for: @mcp.tool(), @app.get/post(), async def with docstrings + types")
        return 0

    print(f"Detected {len(tools)} tool(s):")
    for t in tools:
        decorator_hint = f" [{t.decorator}]" if t.decorator else ""
        params_hint = ", ".join(p.name for p in t.params)
        print(f"  {t.source_file}: {t.name}({params_hint}){decorator_hint}")
    print()

    # 2. Group into skills
    skills = group_into_skills(tools, namespace)

    # 3. Generate skill directories
    skills_dir = output_dir / "skills"
    total_files = 0
    for skill in skills:
        created = generate_skill_directory(skill, skills_dir, source_package)
        total_files += len(created)
        print(f"Generated skill: {skill.namespace}/ ({len(skill.tools)} tools)")
        for f in created:
            print(f"  {f}")
    print()

    # 4. Generate albus.yaml (only if not already present)
    config_path = output_dir / "albus.yaml"
    if config_path.exists():
        print("albus.yaml already exists — skipping (review skills/ and update manually)")
    else:
        config_content = generate_albus_yaml(skills, namespace)
        config_path.write_text(config_content)
        print("Generated albus.yaml")
    print()

    # Summary
    total_tools = sum(len(s.tools) for s in skills)
    print(f"Done. {total_tools} tools → {len(skills)} skill(s) → {total_files} files")
    print()
    print("Next steps:")
    print("  1. Review skills/ — check SKILL.md and tool wrappers")
    print("  2. Review albus.yaml — customize agents and model routing")
    print(f"  3. pip install the source package ({source_package})")
    print("  4. albus server")

    return total_tools


def _infer_namespace(source_dir: Path) -> str:
    """Infer a skill namespace from the source directory.

    Strips common suffixes like _mcp, _api, _server, _tools.
    """
    name = source_dir.name

    # If we're pointing at a src/ subdirectory, use the package name
    if name == "src":
        # Look for a single package directory inside src
        subdirs = [d for d in source_dir.iterdir() if d.is_dir() and not d.name.startswith(".")]
        if len(subdirs) == 1:
            name = subdirs[0].name

    # Strip common suffixes
    for suffix in ("_mcp", "_api", "_server", "_service", "_tools", "_client"):
        if name.endswith(suffix):
            name = name[: -len(suffix)]
            break

    return name.lower().replace("-", "_")


def _infer_source_package(source_dir: Path) -> str:
    """Infer the Python package name from the source directory.

    If source_dir is a src/ layout, returns the package inside it.
    Otherwise returns the directory name itself.
    """
    if source_dir.name == "src":
        subdirs = [d for d in source_dir.iterdir() if d.is_dir() and not d.name.startswith(".")]
        if len(subdirs) == 1:
            return subdirs[0].name

    # Check if source_dir IS a Python package
    if (source_dir / "__init__.py").exists():
        return source_dir.name

    # Check parent is src/
    if source_dir.parent.name == "src":
        return source_dir.name

    return source_dir.name


__all__ = [
    "DetectedParam",
    "DetectedSkill",
    "DetectedTool",
    "generate_albus_yaml",
    "generate_skill_directory",
    "group_into_skills",
    "run_scan",
    "scan_directory",
]
