"""abom scan — statically detect an AI agent's components from a repository.

Produces the `agent`, `components`, and `controls` of an ABOM Composition
Manifest by inspecting dependency manifests, source, prompt files, and MCP
configs. Pure stdlib. Best-effort and conservative: it reports what it
can see, with where it saw it (`detected_from`).
"""
from __future__ import annotations

import hashlib
import json
import logging
import re
from pathlib import Path

try:
    import tomllib  # Python 3.11+
except ModuleNotFoundError:  # Python 3.10
    import tomli as tomllib

log = logging.getLogger("abom.scan")

# --- known signals -----------------------------------------------------------
FRAMEWORKS = {
    "langchain": "LangChain", "langchain-core": "LangChain", "langgraph": "LangGraph",
    "crewai": "CrewAI", "llama-index": "LlamaIndex", "llama_index": "LlamaIndex",
    "autogen": "AutoGen", "pyautogen": "AutoGen", "autogen-agentchat": "AutoGen",
    "semantic-kernel": "Semantic Kernel", "haystack-ai": "Haystack",
    "smolagents": "smolagents", "openai-agents": "OpenAI Agents SDK",
    "google-adk": "Google ADK", "google-genai": "Google GenAI", "mcp": "Model Context Protocol",
    "dspy": "DSPy", "dspy-ai": "DSPy", "pydantic-ai": "Pydantic AI",
}
MODEL_SDKS = {
    "openai": "OpenAI", "anthropic": "Anthropic", "mistralai": "Mistral",
    "cohere": "Cohere", "google-generativeai": "Google", "vertexai": "Google Vertex",
    "boto3": "AWS Bedrock", "ollama": "Ollama", "vllm": "vLLM",
    "huggingface-hub": "Hugging Face", "transformers": "Hugging Face",
    "groq": "Groq", "together": "Together", "replicate": "Replicate",
}
VECTOR_STORES = {
    "chromadb": "Chroma", "pinecone-client": "Pinecone", "pinecone": "Pinecone",
    "weaviate-client": "Weaviate", "qdrant-client": "Qdrant", "faiss-cpu": "FAISS",
    "faiss-gpu": "FAISS", "pgvector": "pgvector", "pymilvus": "Milvus", "lancedb": "LanceDB",
}
GUARDRAILS = {
    "guardrails-ai": "Guardrails AI", "nemoguardrails": "NeMo Guardrails",
    "llm-guard": "LLM Guard", "presidio-analyzer": "Presidio", "rebuff": "Rebuff",
}

# Model identifiers, using [a-z0-9] (NOT \w) so we don't swallow env-var names
# like MISTRAL_API_KEY or O3_MINI_MODEL (underscores break the match).
MODEL_NAME_RE = re.compile(
    r"\b(gpt-4[a-z0-9.\-]*|gpt-3\.5[a-z0-9.\-]*|o[134](?:-[a-z0-9.\-]+)?|chatgpt[a-z0-9.\-]*|"
    r"claude-[a-z0-9.\-]+|gemini-[0-9][a-z0-9.\-]*|mistral-[a-z0-9.\-]+|mixtral-[a-z0-9.\-]+|"
    r"llama-?[23][a-z0-9.\-]*|qwen[0-9][a-z0-9.\-]*|command-r[a-z0-9.\-]*|"
    r"deepseek-[a-z0-9.\-]+|phi-?[34][a-z0-9.\-]*)\b",
    re.IGNORECASE,
)
# Prose/marketing tokens that look like models but aren't real model identifiers.
_MODEL_STOPWORDS = (
    "code", "agi", "api", "integration", "skill", "desktop", "powered", "class",
    "specific", "compatible", "users", "fight", "vs", "ai-fight", "for-pro",
)
# An env-var-shaped token (ALL_CAPS_WITH_UNDERSCORES) is never a model name.
_ENVVAR_RE = re.compile(r"^[A-Z0-9]+(_[A-Z0-9]+)+$")


def _looks_like_real_model(name: str) -> bool:
    """Filter out prose tokens, env-var names, and obvious non-models.

    Real model ids look like ``claude-3-5-sonnet-20241022`` or ``gpt-4o-mini``;
    prose tokens look like ``claude-code`` or ``o3-vs-googles-gemini-2-0``.
    """
    if _ENVVAR_RE.match(name):                      # MISTRAL_API_KEY, O3_MINI_MODEL
        return False
    segs = name.lower().split("-")
    # a hyphen-segment that is a prose stopword word → not a model
    if any(s in _MODEL_STOPWORDS for s in segs):    # claude-code, gemini-agi, …-vs-…
        return False
    # an excessive run of purely-alphabetic words reads like a phrase, not an id
    # (real ids interleave digits: claude / 3 / 5 / sonnet / 20241022)
    alpha_words = [s for s in segs if s.isalpha() and len(s) > 2]
    if len(alpha_words) >= 4:
        return False
    return True
TOOL_RE = re.compile(r"@(?:tool|function_tool|tool\([^)]*\))\s*(?:\n\s*)*def\s+(\w+)|"
                     r"@(?:tool|function_tool)\b[\s\S]{0,80}?def\s+(\w+)")
MCP_CONFIG_NAMES = {"mcp.json", ".mcp.json", "claude_desktop_config.json", "mcp_config.json", ".cursor/mcp.json"}
IGNORE_DIRS = {".git", "node_modules", ".venv", "venv", "env", "__pycache__", "dist", "build",
               ".mypy_cache", ".pytest_cache", ".tox", ".ruff_cache", "site-packages", ".next"}
TEXT_EXT = {".py", ".js", ".ts", ".tsx", ".yaml", ".yml", ".json", ".toml", ".env", ".cfg", ".ini", ".md", ".txt"}
# Model names are only trusted from CODE / config — NOT from prose (.md/.txt),
# where supported-model lists, blog text and docs produce false positives.
CODE_EXT = {".py", ".js", ".ts", ".tsx", ".yaml", ".yml", ".json", ".toml", ".env", ".cfg", ".ini"}
MAX_FILE = 1_000_000


def _read(p: Path) -> str:
    try:
        return p.read_text(encoding="utf-8", errors="ignore")
    except Exception:
        return ""


def _sha256(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def _iter_files(root: Path):
    for p in root.rglob("*"):
        if not p.is_file():
            continue
        if any(part in IGNORE_DIRS for part in p.parts):
            continue
        try:
            if p.stat().st_size > MAX_FILE:
                continue
        except OSError:
            continue
        yield p


def _norm_dep(name: str) -> str:
    return re.split(r"[<>=!~ ;\[]", name.strip(), maxsplit=1)[0].strip().lower().replace("_", "-")


def _dep_version(spec: str) -> str:
    """Extract a readable version from a requirement spec, '' if unpinned.

    'langchain==0.2.1' -> '0.2.1' · 'openai>=1.0' -> '>=1.0' · 'crewai' -> ''
    """
    m = re.search(r"(?:==|>=|<=|~=|!=|>|<)\s*([0-9][\w.\-*]*)", spec)
    if not m:
        return ""
    op = spec[m.start():m.start(1)].strip()
    ver = m.group(1)
    return ver if op == "==" else f"{op}{ver}"


def parse_dependencies(root: Path) -> dict[str, str]:
    """Return {normalized-name: version-or-''}. A pinned version wins over ''."""
    deps: dict[str, str] = {}

    def put(raw: str, version: str = ""):
        name = _norm_dep(raw)
        if not name:
            return
        # keep the most specific (non-empty) version we've seen for this dep
        if name not in deps or (not deps[name] and version):
            deps[name] = version

    for req in list(root.glob("requirements*.txt")) + list(root.glob("**/requirements*.txt")):
        if any(part in IGNORE_DIRS for part in req.parts):
            continue
        for line in _read(req).splitlines():
            line = line.strip()
            if line and not line.startswith(("#", "-")):
                put(line, _dep_version(line))
    pp = root / "pyproject.toml"
    if pp.exists():
        try:
            data = tomllib.loads(_read(pp))
            for d in data.get("project", {}).get("dependencies", []) or []:
                put(d, _dep_version(d))
            poetry = data.get("tool", {}).get("poetry", {}).get("dependencies", {})
            for d, spec in poetry.items():
                put(d, _dep_version(f"{d}{spec}" if isinstance(spec, str) else d))
        except Exception:
            pass
    pkg = root / "package.json"
    if pkg.exists():
        try:
            data = json.loads(_read(pkg))
            for section in ("dependencies", "devDependencies"):
                for k, v in data.get(section, {}).items():
                    put(k, str(v).lstrip("^~") if isinstance(v, str) else "")
        except Exception:
            pass
    deps.pop("", None)
    return deps


def _agent_meta(root: Path) -> dict:
    name, version = root.resolve().name, "0.0.0"
    pp = root / "pyproject.toml"
    if pp.exists():
        try:
            data = tomllib.loads(_read(pp))
            proj = data.get("project", {})
            name = proj.get("name", name)
            version = proj.get("version", version)
        except Exception:
            pass
    pkg = root / "package.json"
    if pkg.exists():
        try:
            data = json.loads(_read(pkg))
            name = data.get("name", name)
            version = data.get("version", version)
        except Exception:
            pass
    return {"name": name, "version": version}


def scan(root: Path) -> dict:
    root = Path(root)
    log.info("scanning %s", root.resolve())
    deps = parse_dependencies(root)
    log.info("parsed %d dependencies", len(deps))
    log.debug("dependencies: %s", ", ".join(sorted(deps)) or "(none)")
    components: list[dict] = []
    seen: set[tuple] = set()

    def add(comp: dict):
        key = (comp["type"], comp.get("name"))
        if key not in seen:
            seen.add(key)
            components.append(comp)
            log.debug("+ %-10s %-28s [%s]", comp["type"], comp.get("name", "?"),
                      comp.get("detected_from", "?"),
                      extra={"event": "component", "comp_type": comp["type"],
                             "comp_name": comp.get("name"), "source": comp.get("detected_from")})

    # 1. from dependencies (carry the pinned version through when present)
    for d in sorted(deps):
        ver = deps.get(d, "")
        vfield = {"version": ver} if ver else {}
        if d in FRAMEWORKS:
            add({"type": "framework", "name": FRAMEWORKS[d], **vfield,
                 "detected_from": f"dependency:{d}"})
        if d in MODEL_SDKS:
            add({"type": "model", "name": f"{MODEL_SDKS[d]} (SDK)", "provider": MODEL_SDKS[d],
                 **vfield, "egress": d not in ("vllm", "ollama", "transformers"),
                 "detected_from": f"dependency:{d}"})
        if d in VECTOR_STORES:
            add({"type": "dataSource", "kind": "vector-store", "name": VECTOR_STORES[d],
                 **vfield, "detected_from": f"dependency:{d}"})
        if d in GUARDRAILS:
            add({"type": "policy", "kind": "guardrail", "name": GUARDRAILS[d],
                 **vfield, "detected_from": f"dependency:{d}"})

    # 2. concrete model names, prompt files, tools, MCP servers — from source.
    # Track WHERE each model was first seen (file + line) so detected_from is
    # explicit, e.g. "gpt_researcher/config.py:42" instead of a vague "source".
    model_locs: dict[str, str] = {}
    for p in _iter_files(root):
        suffix = p.suffix.lower()
        rel = str(p.relative_to(root))
        # prompt files
        if suffix == ".prompt" or "prompt" in p.parent.name.lower() and suffix in (".txt", ".md"):
            text = _read(p)
            add({"type": "prompt", "name": rel, "sha256": _sha256(text),
                 "detected_from": f"file:{rel}"})
        # mcp configs
        if p.name in MCP_CONFIG_NAMES or rel in MCP_CONFIG_NAMES:
            try:
                cfg = json.loads(_read(p))
                for server in (cfg.get("mcpServers") or {}):
                    add({"type": "mcpServer", "name": server, "detected_from": f"file:{rel}"})
            except Exception:
                pass
        if suffix not in TEXT_EXT:
            continue
        text = _read(p)
        # only trust model names from CODE/config files, not prose (docs),
        # and not from tests or pricing/catalog files (they list models the
        # agent merely *knows about*, not models it actually uses).
        rl = rel.lower()
        is_catalog = (
            "test" in rl or "/tests/" in rl or rl.startswith("tests/")
            or "cost" in rl or "pricing" in rl or "model_list" in rl or "models." in rl
        )
        if suffix in CODE_EXT and not is_catalog:
            # find each model name with the line it first appears on
            found: dict[str, int] = {}
            for lineno, line in enumerate(text.splitlines(), start=1):
                for m in MODEL_NAME_RE.findall(line):
                    cand = m if isinstance(m, str) else next(filter(None, m))
                    if _looks_like_real_model(cand) and cand not in found:
                        found[cand] = lineno
            # a single file emitting many model names is a catalog, not usage
            if len(found) <= 8:
                for cand, lineno in found.items():
                    model_locs.setdefault(cand, f"file:{rel}:{lineno}")
            else:
                log.debug("skipping %s — looks like a model catalog (%d names)", rel, len(found))
        if suffix == ".py" and ("@tool" in text or "@function_tool" in text):
            for groups in TOOL_RE.findall(text):
                fn = next((g for g in groups if g), None)
                if fn:
                    add({"type": "tool", "name": fn, "detected_from": f"file:{rel}"})

    # dedupe case-insensitively (keep the first/canonical casing + its location)
    seen_lower: dict[str, str] = {}
    for name in sorted(model_locs):
        seen_lower.setdefault(name.lower(), name)
    for name in sorted(seen_lower.values()):
        provider = "external" if not name.lower().startswith(("llama", "qwen", "mistral", "mixtral", "phi", "deepseek")) else "open-weight"
        add({"type": "model", "name": name, "provenance": "referenced in source",
             "egress": provider == "external",
             "detected_from": model_locs.get(name, "source")})

    log.info("detected %d components", len(components),
             extra={"event": "scan_done", "components": len(components)})
    return {
        "agent": _agent_meta(root),
        "components": components,
        "controls": {"scan": "static", "scanned_path": str(root.resolve().name),
                     "egress": "unknown (static scan)"},
    }
