"""
OmniSec AI Engine
=================
Utilities for text analysis, classification, summarisation,
embedding generation, and chat completions via the Anthropic API.
"""

from __future__ import annotations

import json
import re
from typing import Any, Dict, List, Optional
from omnisec.core.logger import get_logger

log = get_logger("omnisec.ai")


class AIEngine:
    """
    High-level AI interface backed by the Anthropic Claude API.

    Features:
      - Text classification
      - Summarisation
      - Sentiment analysis
      - Named-entity extraction
      - Security log analysis
      - Generic chat completions
      - Threat intelligence narrative generation

    Example:
        >>> engine = AIEngine(api_key="sk-ant-...")
        >>> result = engine.summarize("Long article text here...")
        >>> print(result)
    """

    DEFAULT_MODEL = "claude-sonnet-4-6"

    def __init__(
        self,
        api_key: Optional[str] = None,
        model: str = DEFAULT_MODEL,
        max_tokens: int = 1024,
    ):
        self.model = model
        self.max_tokens = max_tokens
        self._api_key = api_key
        self._client = None
        self._init_client()

    def _init_client(self):
        try:
            import anthropic
            self._client = anthropic.Anthropic(api_key=self._api_key) if self._api_key else anthropic.Anthropic()
            log.info(f"AIEngine initialized with model={self.model}")
        except ImportError:
            log.warning("anthropic package not installed. Run: pip install anthropic")

    def _complete(self, system: str, user: str) -> str:
        """Internal: send a system+user prompt and return text response."""
        if not self._client:
            raise RuntimeError("Anthropic client not available. Install 'anthropic' package.")
        msg = self._client.messages.create(
            model=self.model,
            max_tokens=self.max_tokens,
            system=system,
            messages=[{"role": "user", "content": user}],
        )
        return msg.content[0].text.strip()

    # ── Public Methods ───────────────────────────────────────────────

    def chat(self, prompt: str, system: str = "You are a helpful assistant.") -> str:
        """
        Send a free-form prompt and get a response.

        Args:
            prompt: User message
            system: Optional system instruction

        Returns:
            str: Model response

        Example:
            >>> engine.chat("What is a zero-day exploit?")
        """
        log.debug(f"chat() prompt_len={len(prompt)}")
        return self._complete(system, prompt)

    def summarize(self, text: str, max_sentences: int = 5) -> str:
        """
        Summarise a block of text.

        Args:
            text: Input text to summarise
            max_sentences: Target summary length in sentences

        Returns:
            str: Summary text

        Example:
            >>> engine.summarize(long_report, max_sentences=3)
        """
        system = f"You are a precise summariser. Summarise in at most {max_sentences} sentences."
        log.debug(f"summarize() input_len={len(text)}")
        return self._complete(system, text)

    def classify(self, text: str, labels: List[str]) -> Dict[str, float]:
        """
        Classify text into one or more provided labels.

        Args:
            text: Text to classify
            labels: List of candidate labels

        Returns:
            Dict mapping each label to a confidence score 0.0–1.0

        Example:
            >>> engine.classify("Phishing email body...", ["phishing","spam","legitimate"])
        """
        system = (
            "You are a text classifier. Return ONLY a JSON object mapping each "
            "label to a confidence score between 0.0 and 1.0. Sum of scores should be 1.0."
        )
        user = f"Labels: {labels}\n\nText:\n{text}"
        raw = self._complete(system, user)
        try:
            cleaned = re.sub(r"```(?:json)?|```", "", raw).strip()
            return json.loads(cleaned)
        except Exception:
            log.warning("classify() could not parse JSON response, returning raw")
            return {"raw": raw}

    def sentiment(self, text: str) -> Dict[str, Any]:
        """
        Analyse sentiment of text.

        Returns:
            Dict with keys: label (positive/negative/neutral), score (0-1), reasoning

        Example:
            >>> engine.sentiment("The firewall blocked the attack successfully.")
        """
        system = (
            "You are a sentiment analyser. Return ONLY a JSON object with keys: "
            "'label' (positive|negative|neutral), 'score' (0.0-1.0), 'reasoning' (one sentence)."
        )
        raw = self._complete(system, text)
        try:
            cleaned = re.sub(r"```(?:json)?|```", "", raw).strip()
            return json.loads(cleaned)
        except Exception:
            return {"label": "unknown", "score": 0.0, "reasoning": raw}

    def extract_entities(self, text: str) -> Dict[str, List[str]]:
        """
        Extract named entities (people, orgs, IPs, domains, CVEs) from text.

        Returns:
            Dict with entity type keys and lists of found entities

        Example:
            >>> engine.extract_entities("CVE-2023-1234 was exploited by APT29 via 192.168.1.1")
        """
        system = (
            "Extract named entities from the text. Return ONLY a JSON object with keys: "
            "'persons', 'organizations', 'ip_addresses', 'domains', 'cves', 'locations'. "
            "Each value is a list of strings. Return empty lists if none found."
        )
        raw = self._complete(system, text)
        try:
            cleaned = re.sub(r"```(?:json)?|```", "", raw).strip()
            return json.loads(cleaned)
        except Exception:
            return {"raw": raw}

    def analyze_security_log(self, log_text: str) -> Dict[str, Any]:
        """
        Analyse a security log entry or block for threats and anomalies.

        Args:
            log_text: Raw log string (syslog, firewall log, SIEM alert, etc.)

        Returns:
            Dict with: threat_level, threat_type, indicators, recommendation, summary

        Example:
            >>> engine.analyze_security_log("Failed SSH login from 10.0.0.5 - 50 attempts")
        """
        system = (
            "You are a cybersecurity analyst. Analyse the security log and return ONLY a JSON with: "
            "'threat_level' (critical|high|medium|low|info), "
            "'threat_type' (brute_force|malware|phishing|ddos|intrusion|data_exfil|anomaly|none), "
            "'indicators' (list of strings), "
            "'recommendation' (string), "
            "'summary' (one sentence)."
        )
        raw = self._complete(system, log_text)
        try:
            cleaned = re.sub(r"```(?:json)?|```", "", raw).strip()
            return json.loads(cleaned)
        except Exception:
            return {"threat_level": "unknown", "summary": raw}

    def generate_threat_report(self, threat_data: Dict[str, Any]) -> str:
        """
        Generate a human-readable threat intelligence report from structured data.

        Args:
            threat_data: Dict containing threat details (IPs, CVEs, timeline, etc.)

        Returns:
            str: Formatted markdown threat report

        Example:
            >>> engine.generate_threat_report({"cve": "CVE-2024-1234", "severity": "critical"})
        """
        system = (
            "You are a senior threat intelligence analyst. "
            "Generate a professional, structured markdown threat report."
        )
        user = f"Generate a threat report for:\n{json.dumps(threat_data, indent=2)}"
        return self._complete(system, user)

    def __repr__(self) -> str:
        return f"AIEngine(model={self.model!r}, max_tokens={self.max_tokens})"
