"""AYN Antivirus — AI-Powered Threat Analyzer.

Uses Claude to analyze suspicious files and filter false positives.
Each detection from heuristic/signature scanners is verified by AI
before being reported as a real threat.
"""

from __future__ import annotations

import json
import logging
import os
import platform
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

SYSTEM_PROMPT = """Linux VPS antivirus analyst. {environment}
Normal: pip/npm scripts in /usr/local/bin, Docker hex IDs, cron jobs (fstrim/certbot/logrotate), high-entropy archives, curl/wget in deploy scripts, recently-modified files after apt/pip.
Reply ONLY JSON: {{"verdict":"threat"|"safe"|"suspicious","confidence":0-100,"reason":"short","recommended_action":"quarantine"|"delete"|"ignore"|"monitor"}}"""

ANALYSIS_PROMPT = """FILE:{file_path} DETECT:{threat_name}({threat_type}) SEV:{severity} DET:{detector} CONF:{original_confidence}% SIZE:{file_size} PERM:{permissions} OWN:{owner} MOD:{mtime}
PREVIEW:
{content_preview}
JSON verdict:"""


@dataclass
class AIVerdict:
    """Result of AI analysis on a detection."""
    verdict: str          # threat, safe, suspicious
    confidence: int       # 0-100
    reason: str
    recommended_action: str  # quarantine, delete, ignore, monitor
    raw_response: str = ""

    @property
    def is_threat(self) -> bool:
        return self.verdict == "threat"

    @property
    def is_safe(self) -> bool:
        return self.verdict == "safe"


class AIAnalyzer:
    """AI-powered threat analysis using Claude."""

    def __init__(self, api_key: Optional[str] = None, model: str = "claude-sonnet-4-20250514"):
        self._api_key = api_key or os.environ.get("ANTHROPIC_API_KEY", "") or self._load_key_from_env_file()
        self._model = model
        self._client = None
        self._environment = self._detect_environment()

    @staticmethod
    def _load_key_from_env_file() -> str:
        for p in ["/opt/ayn-antivirus/.env", Path.home() / ".ayn-antivirus" / ".env"]:
            try:
                for line in Path(p).read_text().splitlines():
                    line = line.strip()
                    if line.startswith("ANTHROPIC_API_KEY=") and not line.endswith("="):
                        return line.split("=", 1)[1].strip().strip("'\"")
            except Exception:
                pass
        return ""

    @property
    def available(self) -> bool:
        return bool(self._api_key)

    def _get_client(self):
        if not self._client:
            try:
                import anthropic
                self._client = anthropic.Anthropic(api_key=self._api_key)
            except Exception as exc:
                logger.error("Failed to init Anthropic client: %s", exc)
                return None
        return self._client

    @staticmethod
    def _detect_environment() -> str:
        """Gather environment context for the AI."""
        import shutil
        parts = [
            f"OS: {platform.system()} {platform.release()}",
            f"Hostname: {platform.node()}",
            f"Arch: {platform.machine()}",
        ]
        if shutil.which("incus"):
            parts.append("Container runtime: Incus/LXC (containers run Docker inside)")
        if shutil.which("docker"):
            parts.append("Docker: available")
        if Path("/etc/dokploy").exists() or shutil.which("dokploy"):
            parts.append("Platform: Dokploy (Docker deployment platform)")

        # Check if we're inside a container
        if Path("/run/host/container-manager").exists():
            parts.append("Running inside: managed container")
        return "\n".join(parts)

    def _get_file_context(self, file_path: str) -> Dict[str, Any]:
        """Gather file metadata and content preview."""
        p = Path(file_path)
        ctx = {
            "file_size": 0,
            "permissions": "",
            "owner": "",
            "mtime": "",
            "content_preview": "[file not readable]",
        }
        try:
            st = p.stat()
            ctx["file_size"] = st.st_size
            ctx["permissions"] = oct(st.st_mode)[-4:]
            ctx["mtime"] = str(st.st_mtime)
            try:
                import pwd
                ctx["owner"] = pwd.getpwuid(st.st_uid).pw_name
            except Exception:
                ctx["owner"] = str(st.st_uid)
        except OSError:
            pass

        try:
            with open(file_path, "rb") as f:
                raw = f.read(512)
            # Try text decode, fall back to hex
            try:
                ctx["content_preview"] = raw.decode("utf-8", errors="replace")
            except Exception:
                ctx["content_preview"] = raw.hex()[:512]
        except Exception:
            pass

        return ctx

    def analyze(
        self,
        file_path: str,
        threat_name: str,
        threat_type: str,
        severity: str,
        detector: str,
        confidence: int = 50,
    ) -> AIVerdict:
        """Analyze a single detection with AI."""
        if not self.available:
            # No API key — pass through as-is
            return AIVerdict(
                verdict="suspicious",
                confidence=confidence,
                reason="AI analysis unavailable (no API key)",
                recommended_action="quarantine",
            )

        client = self._get_client()
        if not client:
            return AIVerdict(
                verdict="suspicious",
                confidence=confidence,
                reason="AI client init failed",
                recommended_action="quarantine",
            )

        ctx = self._get_file_context(file_path)

        # Sanitize content preview to avoid format string issues
        preview = ctx.get("content_preview", "")
        if len(preview) > 500:
            preview = preview[:500] + "..."
        # Replace curly braces to avoid format() issues
        preview = preview.replace("{", "{{").replace("}", "}}")

        user_msg = ANALYSIS_PROMPT.format(
            file_path=file_path,
            threat_name=threat_name,
            threat_type=threat_type,
            severity=severity,
            detector=detector,
            original_confidence=confidence,
            file_size=ctx.get("file_size", 0),
            permissions=ctx.get("permissions", ""),
            owner=ctx.get("owner", ""),
            mtime=ctx.get("mtime", ""),
            content_preview=preview,
        )

        text = ""
        try:
            response = client.messages.create(
                model=self._model,
                max_tokens=150,
                system=SYSTEM_PROMPT.format(environment=self._environment),
                messages=[{"role": "user", "content": user_msg}],
            )
            text = response.content[0].text.strip()

            # Parse JSON from response (handle markdown code blocks)
            if "```" in text:
                parts = text.split("```")
                for part in parts[1:]:
                    cleaned = part.strip()
                    if cleaned.startswith("json"):
                        cleaned = cleaned[4:].strip()
                    if cleaned.startswith("{"):
                        text = cleaned
                        break

            # Find the JSON object in the response
            start = text.find("{")
            end = text.rfind("}") + 1
            if start >= 0 and end > start:
                text = text[start:end]

            data = json.loads(text)
            return AIVerdict(
                verdict=data.get("verdict", "suspicious"),
                confidence=data.get("confidence", 50),
                reason=data.get("reason", ""),
                recommended_action=data.get("recommended_action", "quarantine"),
                raw_response=text,
            )
        except json.JSONDecodeError as exc:
            logger.warning("AI returned non-JSON: %s — raw: %s", exc, text[:200])
            return AIVerdict(
                verdict="suspicious",
                confidence=confidence,
                reason=f"AI parse error: {text[:100]}",
                recommended_action="quarantine",
                raw_response=text,
            )
        except Exception as exc:
            logger.error("AI analysis failed: %s", exc)
            return AIVerdict(
                verdict="suspicious",
                confidence=confidence,
                reason=f"AI error: {exc}",
                recommended_action="quarantine",
            )

    def analyze_batch(
        self,
        detections: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """Analyze a batch of detections. Returns enriched detections with AI verdicts.

        Each detection dict should have: file_path, threat_name, threat_type, severity, detector
        """
        results = []
        for d in detections:
            verdict = self.analyze(
                file_path=d.get("file_path", ""),
                threat_name=d.get("threat_name", ""),
                threat_type=d.get("threat_type", ""),
                severity=d.get("severity", "MEDIUM"),
                detector=d.get("detector", ""),
                confidence=d.get("confidence", 50),
            )
            enriched = dict(d)
            enriched["ai_verdict"] = verdict.verdict
            enriched["ai_confidence"] = verdict.confidence
            enriched["ai_reason"] = verdict.reason
            enriched["ai_action"] = verdict.recommended_action
            results.append(enriched)
        return results