"""
Pattern Cache — syncs known failure patterns + fixes from platform.

Periodically pulls patterns from GET /v1/remediation/patterns
and updates the FastDetector's cache for instant <5ms matching.

Background refresh runs every 5 minutes. Initial sync on SDK init.
"""

import asyncio
import hashlib
import logging
import time
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    from ..judge.kg_client import KnowledgeGraphClient
    from ..judge.selector import JudgeSelector
    from ..judge.trajectory_tracker import TrajectoryTracker
    from .fast_detector import FastDetector

logger = logging.getLogger("aigie.runtime.pattern_cache")


class PatternCache:
    """Syncs remediation patterns from the Aigie platform.

    Thread-safe: the fast_detector's cache is updated via atomic dict swap.
    """

    def __init__(
        self,
        api_url: str,
        api_key: str | None = None,
        refresh_interval_sec: float = 300.0,  # 5 minutes
        fast_detector: Optional["FastDetector"] = None,
        judge_selector: Optional["JudgeSelector"] = None,
        kg_client: Optional["KnowledgeGraphClient"] = None,
        trajectory_tracker: Optional["TrajectoryTracker"] = None,
    ):
        self._api_url = api_url.rstrip("/")
        self._api_key = api_key
        self._refresh_interval = refresh_interval_sec
        self._fast_detector = fast_detector
        self._judge_selector = judge_selector
        self._kg_client = kg_client
        self._trajectory_tracker = trajectory_tracker
        self._accuracy_map: dict[str, float] = {}
        self._refresh_task: asyncio.Task | None = None
        self._last_sync: float = 0
        self._pattern_count: int = 0
        self._platform_judge_count: int = 0
        self._enrichment_data: dict[str, dict] = {}  # agent_name/workflow_id → enrichment
        self._http_client = None  # Persistent HTTP client (lazily initialized)

    async def initial_sync(self) -> int:
        """Perform initial pattern sync. Returns number of patterns loaded."""
        count = 0
        try:
            patterns, rules, judges = await asyncio.gather(
                self._fetch_patterns(),
                self._fetch_custom_rules(),
                self._fetch_platform_judges(),
                return_exceptions=True,
            )
            if isinstance(patterns, list) and patterns:
                self._apply_patterns(patterns)
                count = len(patterns)
            if isinstance(rules, list) and rules and self._fast_detector:
                self._fast_detector.update_custom_rules(rules)
                logger.debug(f"Loaded {len(rules)} custom detection rules")
            if isinstance(judges, list) and judges and self._judge_selector:
                self._judge_selector.update_platform_judges(judges)
                self._platform_judge_count = len(judges)
                logger.debug(f"Loaded {len(judges)} platform judges")
            # Sync KG fixes
            if self._kg_client:
                try:
                    fixes = await self._fetch_top_fixes()
                    if fixes:
                        self._kg_client.update_cache(fixes)
                        logger.debug(f"KG cache loaded {len(fixes)} fixes")
                except Exception as e:
                    logger.debug(f"KG sync failed (non-fatal): {e}")
            # Sync judge accuracy
            try:
                accuracy_map = await self._fetch_accuracy_map()
                if accuracy_map:
                    self._accuracy_map = accuracy_map
                    logger.debug(f"Loaded accuracy for {len(accuracy_map)} judge domains")
            except Exception as e:
                logger.debug(f"Accuracy sync failed (non-fatal): {e}")
            # Sync trajectory patterns
            if self._trajectory_tracker:
                try:
                    patterns = await self._fetch_trajectory_patterns()
                    if patterns:
                        self._trajectory_tracker.update_patterns(patterns)
                        logger.debug(f"Loaded {len(patterns)} trajectory patterns")
                except Exception as e:
                    logger.debug(f"Trajectory sync failed (non-fatal): {e}")
            self._last_sync = time.monotonic()
            if count:
                logger.info(f"Pattern cache loaded {count} patterns")
        except Exception as e:
            logger.debug(f"Initial pattern sync failed (non-fatal): {e}")
        return count

    def start_background_refresh(self) -> None:
        """Start background refresh task."""
        if self._refresh_task is not None:
            return
        try:
            loop = asyncio.get_running_loop()
            self._refresh_task = loop.create_task(self._refresh_loop())
        except RuntimeError:
            logger.debug("No event loop — pattern refresh disabled")

    async def force_refresh(self) -> int:
        """Force an immediate pattern refresh (e.g. after backend push notification).

        Returns:
            Number of patterns loaded.
        """
        try:
            patterns, rules, judges = await asyncio.gather(
                self._fetch_patterns(),
                self._fetch_custom_rules(),
                self._fetch_platform_judges(),
                return_exceptions=True,
            )
            count = 0
            if isinstance(patterns, list) and patterns:
                self._apply_patterns(patterns)
                count = len(patterns)
            if isinstance(rules, list) and rules and self._fast_detector:
                self._fast_detector.update_custom_rules(rules)
            if isinstance(judges, list) and judges and self._judge_selector:
                self._judge_selector.update_platform_judges(judges)
                self._platform_judge_count = len(judges)
            # Refresh KG fixes
            if self._kg_client:
                try:
                    fixes = await self._fetch_top_fixes()
                    if fixes:
                        self._kg_client.update_cache(fixes)
                except Exception as e:
                    logger.debug(f"KG force-refresh failed (non-fatal): {e}")
            # Refresh accuracy
            try:
                accuracy_map = await self._fetch_accuracy_map()
                if accuracy_map:
                    self._accuracy_map = accuracy_map
            except Exception as e:
                logger.debug(f"Accuracy force-refresh failed (non-fatal): {e}")
            # Refresh trajectory patterns
            if self._trajectory_tracker:
                try:
                    patterns = await self._fetch_trajectory_patterns()
                    if patterns:
                        self._trajectory_tracker.update_patterns(patterns)
                except Exception as e:
                    logger.debug(f"Trajectory force-refresh failed (non-fatal): {e}")
            self._last_sync = time.monotonic()
            if count:
                logger.info(f"Pattern cache force-refreshed: {count} patterns")
            return count
        except Exception as e:
            logger.debug(f"Force refresh failed: {e}")
            return 0

    async def _refresh_loop(self) -> None:
        """Background refresh loop."""
        while True:
            try:
                await asyncio.sleep(self._refresh_interval)
                patterns, rules, judges = await asyncio.gather(
                    self._fetch_patterns(),
                    self._fetch_custom_rules(),
                    self._fetch_platform_judges(),
                    return_exceptions=True,
                )
                if isinstance(patterns, list) and patterns:
                    self._apply_patterns(patterns)
                if isinstance(rules, list) and rules and self._fast_detector:
                    self._fast_detector.update_custom_rules(rules)
                if isinstance(judges, list) and judges and self._judge_selector:
                    self._judge_selector.update_platform_judges(judges)
                    self._platform_judge_count = len(judges)
                # Refresh KG fixes
                if self._kg_client:
                    try:
                        fixes = await self._fetch_top_fixes()
                        if fixes:
                            self._kg_client.update_cache(fixes)
                    except Exception as e:
                        logger.debug(f"KG refresh failed (non-fatal): {e}")
                # Refresh accuracy
                try:
                    accuracy_map = await self._fetch_accuracy_map()
                    if accuracy_map:
                        self._accuracy_map = accuracy_map
                except Exception as e:
                    logger.debug(f"Accuracy refresh failed (non-fatal): {e}")
                # Refresh trajectory patterns
                if self._trajectory_tracker:
                    try:
                        patterns = await self._fetch_trajectory_patterns()
                        if patterns:
                            self._trajectory_tracker.update_patterns(patterns)
                    except Exception as e:
                        logger.debug(f"Trajectory refresh failed (non-fatal): {e}")
                self._last_sync = time.monotonic()
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.debug(f"Pattern refresh failed: {e}")

    async def _get_http_client(self):
        """Get or create persistent HTTP client."""
        if self._http_client is None:
            import httpx

            headers = {"Content-Type": "application/json"}
            if self._api_key:
                headers["X-API-Key"] = self._api_key
            self._http_client = httpx.AsyncClient(
                timeout=10.0,
                headers=headers,
            )
        return self._http_client

    async def _fetch_patterns(self) -> List[Dict[str, Any]]:
        """Fetch patterns from platform API."""
        client = await self._get_http_client()
        url = f"{self._api_url}/v1/remediation/patterns"
        resp = await client.get(url)
        if resp.status_code == 200:
            data = resp.json()
            return data if isinstance(data, list) else data.get("patterns", [])
        return []

    async def _fetch_custom_rules(self) -> List[Dict[str, Any]]:
        """Fetch custom quality rules from platform API."""
        try:
            client = await self._get_http_client()
            url = f"{self._api_url}/v1/detection/rules"
            resp = await client.get(url)
            if resp.status_code == 200:
                data = resp.json()
                return data if isinstance(data, list) else data.get("rules", [])
        except Exception as e:
            logger.debug(f"Custom rules fetch failed (non-fatal): {e}")
        return []

    async def _fetch_platform_judges(self) -> List[Dict[str, Any]]:
        """Fetch predefined judges from Kytte platform.

        Returns list of judge dicts with: judge_id, name, category,
        template_id, signal_triggers, span_types, skills, priority.
        """
        try:
            client = await self._get_http_client()
            url = f"{self._api_url}/v1/evals/judges"
            resp = await client.get(url)
            if resp.status_code == 200:
                data = resp.json()
                return data if isinstance(data, list) else data.get("judges", [])
        except Exception as e:
            logger.debug(f"Platform judges fetch failed (non-fatal): {e}")
        return []

    async def _fetch_trajectory_patterns(self) -> List[Dict[str, Any]]:
        """Fetch active trajectory patterns from platform API."""
        try:
            client = await self._get_http_client()
            url = f"{self._api_url}/v1/trajectory/patterns"
            resp = await client.get(url)
            if resp.status_code == 200:
                data = resp.json()
                return data.get("items", [])
        except Exception as e:
            logger.debug(f"Trajectory patterns fetch failed (non-fatal): {e}")
        return []

    async def _fetch_accuracy_map(self) -> Dict[str, float]:
        """Fetch judge accuracy map from platform API."""
        try:
            client = await self._get_http_client()
            url = f"{self._api_url}/v1/judge-accuracy/map"
            resp = await client.get(url)
            if resp.status_code == 200:
                return resp.json()
        except Exception as e:
            logger.debug(f"Accuracy map fetch failed (non-fatal): {e}")
        return {}

    async def _fetch_top_fixes(self) -> Dict[str, Dict[str, Any]]:
        """Fetch top KG fixes from platform and return as cache dict."""
        try:
            client = await self._get_http_client()
            url = f"{self._api_url}/v1/knowledge-graph/stats"
            resp = await client.get(url)
            if resp.status_code != 200:
                return {}
            # Fetch fixes for known domains
            fixes: Dict[str, Dict[str, Any]] = {}
            for domain in ("hallucination", "drift", "safety", "quality", "tool_usage", "loop"):
                for root_cause in self._get_domain_root_causes(domain):
                    fix_resp = await client.get(
                        f"{self._api_url}/v1/knowledge-graph/fix",
                        params={"issue_domain": domain, "root_cause": root_cause},
                    )
                    if fix_resp.status_code == 200:
                        data = fix_resp.json().get("fix")
                        if data:
                            fixes[f"{domain}|{root_cause}"] = data
            return fixes
        except Exception as e:
            logger.debug(f"KG fixes fetch failed (non-fatal): {e}")
            return {}

    @staticmethod
    def _get_domain_root_causes(domain: str) -> List[str]:
        """Known root causes per domain for KG queries."""
        return {
            "hallucination": [
                "ungrounded_claim",
                "tool_result_mismatch",
                "fabricated_reference",
                "overconfident_assertion",
            ],
            "drift": ["task_deviation", "topic_shift", "role_deviation"],
            "safety": ["harmful_content", "pii_leak", "injection_success", "policy_violation"],
            "quality": [
                "incomplete_response",
                "incoherent_logic",
                "irrelevant_content",
                "vague_response",
            ],
            "tool_usage": [
                "wrong_tool",
                "bad_parameters",
                "misinterpreted_result",
                "unhandled_tool_error",
            ],
            "loop": ["stuck_loop", "circular_reasoning", "tool_retry_loop"],
        }.get(domain, [])

    async def _fetch_enrichment(
        self, agent_name: str | None = None, workflow_id: str | None = None
    ) -> dict:
        """Fetch enrichment context from platform API."""
        try:
            client = await self._get_http_client()
            params = {}
            if agent_name:
                params["agent_name"] = agent_name
            if workflow_id:
                params["workflow_id"] = workflow_id
            url = f"{self._api_url}/v1/enrichment/context"
            resp = await client.get(url, params=params)
            if resp.status_code == 200:
                return resp.json()
        except Exception as e:
            logger.debug(f"Enrichment fetch failed (non-fatal): {e}")
        return {}

    def get_enrichment_context(
        self, agent_name: str | None = None, workflow_id: str | None = None
    ) -> dict:
        """Get cached enrichment context. Returns empty dict if not available."""
        key = f"{agent_name or ''}:{workflow_id or ''}"
        data = dict(self._enrichment_data.get(key, {}))
        if self._accuracy_map:
            data["judge_accuracy"] = self._accuracy_map
        return data

    def update_enrichment_context(
        self, agent_name: str | None, workflow_id: str | None, data: dict
    ) -> None:
        """Update cached enrichment context."""
        key = f"{agent_name or ''}:{workflow_id or ''}"
        self._enrichment_data[key] = data

    def _apply_patterns(self, patterns: List[Dict[str, Any]]) -> None:
        """Convert platform patterns to fast-lookup cache and apply."""
        cache: Dict[str, Dict[str, Any]] = {}
        for p in patterns:
            # Index by error type
            error_type = p.get("error_type", "")
            if error_type:
                cache[error_type] = p

            # Also index by error signature hash for fuzzy matching
            heuristic = p.get("judge_heuristic", "")
            if heuristic:
                heuristic_str = str(heuristic) if not isinstance(heuristic, str) else heuristic
                sig = hashlib.md5(
                    heuristic_str[:200].lower().encode(), usedforsecurity=False
                ).hexdigest()[:16]
                cache[sig] = p

            # Index by pattern name
            name = p.get("name", "")
            if name:
                cache[name] = p

        self._pattern_count = len(patterns)
        if self._fast_detector:
            self._fast_detector.update_pattern_cache(cache)

    async def stop(self) -> None:
        """Stop background refresh."""
        if self._refresh_task:
            self._refresh_task.cancel()
            try:
                await self._refresh_task
            except asyncio.CancelledError:
                pass
            self._refresh_task = None

        if self._http_client:
            await self._http_client.aclose()
            self._http_client = None

    def get_stats(self) -> Dict[str, Any]:
        return {
            "pattern_count": self._pattern_count,
            "platform_judge_count": self._platform_judge_count,
            "last_sync_ago_sec": time.monotonic() - self._last_sync if self._last_sync else None,
            "refresh_interval_sec": self._refresh_interval,
        }
