"""Layer 2 — AI analysis using the Anthropic Claude API."""
from __future__ import annotations

import json
import logging
import os
from pathlib import Path
from typing import Any

from pyguard.layers.base import BaseLayer
from pyguard.models import Finding, PackageInfo, Severity

logger = logging.getLogger(__name__)

# Files analysed first (highest attack surface)
_PRIORITY_NAMES = {"setup.py", "setup.cfg", "pyproject.toml", "__init__.py"}

# Rough token budget per request (leaves head-room in 200 k context window)
_MAX_SOURCE_CHARS = 120_000

_SYSTEM_PROMPT = """\
You are a Python supply chain security expert. Analyse the provided package source code \
for ACTIVE malicious behaviour — not theoretical vulnerabilities.

Focus on:
1. Credential theft: reading ~/.ssh/, ~/.aws/, ~/.kube/, env vars with key/token/secret.
2. Data exfiltration: HTTP POST to external domains during install or import.
3. Persistence: modifying shell configs, cron jobs, startup items.
4. Obfuscation: base64-encoded payloads, exec(compile(…)), __import__ tricks.
5. Install hooks: setup.py cmdclass overrides, distutils hooks.
6. Privilege escalation: sudo, setuid, modifying system files.

Respond ONLY with valid JSON matching this schema (no markdown, no prose):
{
  "findings": [
    {
      "severity": "critical|high|medium|low|info",
      "title": "Short title (≤80 chars)",
      "description": "What it does and why it is dangerous",
      "file": "relative/path/to/file.py",
      "line_hint": 42,
      "code_snippet": "1–3 relevant lines",
      "confidence": 0.95,
      "category": "credential_theft|exfiltration|obfuscation|persistence|execution|other"
    }
  ],
  "overall_verdict": "malicious|suspicious|clean",
  "reasoning": "One-sentence explanation"
}
If there are no findings return {"findings": [], "overall_verdict": "clean", "reasoning": "…"}.
"""

_SEVERITY_MAP: dict[str, Severity] = {
    "critical": Severity.CRITICAL,
    "high": Severity.HIGH,
    "medium": Severity.MEDIUM,
    "low": Severity.LOW,
    "info": Severity.INFO,
}


def _priority_sort_key(file_path: str) -> int:
    return 0 if Path(file_path).name in _PRIORITY_NAMES else 1


def _build_prompt(package: PackageInfo) -> str:
    """Assemble the user message: package identity + source files."""
    lines: list[str] = [
        f"Package: {package.name}=={package.version}",
        "",
        "Source files (in priority order):",
        "",
    ]
    chars_used = 0
    for file_path, content in sorted(package.files.items(), key=lambda kv: _priority_sort_key(kv[0])):
        if not file_path.endswith(".py"):
            continue
        header = f"### {file_path}\n"
        block = f"```python\n{content}\n```\n\n"
        chunk = header + block
        if chars_used + len(chunk) > _MAX_SOURCE_CHARS:
            lines.append(f"### {file_path}\n[truncated — file budget exceeded]\n")
            break
        lines.append(chunk)
        chars_used += len(chunk)
    return "\n".join(lines)


def _parse_response(
    package: PackageInfo, raw: str
) -> tuple[list[Finding], str | None, float | None]:
    """Parse a Claude JSON response into findings + verdict."""
    try:
        data: dict[str, Any] = json.loads(raw)
    except json.JSONDecodeError:
        logger.warning("AI layer: could not parse JSON response")
        return [], None, None

    findings: list[Finding] = []
    for item in data.get("findings", []):
        sev_str = str(item.get("severity", "medium")).lower()
        severity = _SEVERITY_MAP.get(sev_str, Severity.MEDIUM)
        confidence = float(item.get("confidence", 0.7))
        findings.append(
            Finding(
                severity=severity,
                title=str(item.get("title", "AI finding"))[:120],
                description=str(item.get("description", "")),
                layer="ai",
                package=package.name,
                version=package.version,
                evidence={
                    "code_snippet": str(item.get("code_snippet", "")),
                    "ai_verdict": str(data.get("overall_verdict", "")),
                },
                confidence=min(max(confidence, 0.0), 1.0),
                file=item.get("file") or None,
                line=int(item["line_hint"]) if item.get("line_hint") else None,
                category=str(item.get("category", "general")),
            )
        )

    verdict: str | None = data.get("overall_verdict") or None
    # Overall confidence = average of individual findings, or 0.9 for explicit verdicts
    avg_conf: float | None = (
        sum(f.confidence for f in findings) / len(findings) if findings else (0.9 if verdict else None)
    )
    return findings, verdict, avg_conf


class AIAnalysisLayer(BaseLayer):
    name = "ai"

    def __init__(self, api_key: str | None = None, model: str = "claude-sonnet-4-6") -> None:
        self._api_key = api_key or os.environ.get("ANTHROPIC_API_KEY", "")
        self._model = model

    async def analyze(self, package: PackageInfo) -> list[Finding]:
        if not self._api_key:
            logger.warning(
                "AI layer disabled: set ANTHROPIC_API_KEY to enable Claude analysis."
            )
            return []

        if not package.files:
            return []

        try:
            import anthropic  # noqa: PLC0415
        except ImportError:
            logger.warning("AI layer disabled: install 'anthropic' package to enable.")
            return []

        prompt = _build_prompt(package)
        client = anthropic.AsyncAnthropic(api_key=self._api_key)

        try:
            message = await client.messages.create(
                model=self._model,
                max_tokens=4096,
                system=_SYSTEM_PROMPT,
                messages=[{"role": "user", "content": prompt}],
            )
        except Exception as exc:  # noqa: BLE001
            logger.warning("AI layer API call failed: %s", exc)
            return []

        raw = message.content[0].text if message.content else ""
        findings, verdict, avg_conf = _parse_response(package, raw)

        # Store verdict on the package scan result via a synthetic INFO finding
        if verdict and verdict != "clean":
            findings.append(
                Finding(
                    severity=Severity.INFO,
                    title=f"AI overall verdict: {verdict}",
                    description=str(
                        json.loads(raw).get("reasoning", "")
                        if raw
                        else ""
                    ),
                    layer="ai",
                    package=package.name,
                    version=package.version,
                    evidence={"verdict": verdict},
                    confidence=avg_conf or 0.8,
                    category="ai_verdict",
                )
            )

        return findings
