"""Diagnoser — rule-based diagnosis with AI fallback.

Analyzes failures to determine root cause and suggest fixes.
Pipeline:
1. Try rule-based fixes (fast, deterministic)
2. Try context-aware heuristics (medium confidence)
3. Fall back to AI diagnosis if available (slower, more creative)
4. Return default retry as last resort
"""

import json
from typing import Any

from xyglue.heal.models import Failure, Diagnosis


class Diagnoser:
    """Diagnoses failures and suggests fixes.

    Uses a tiered approach:
    1. Rule-based: deterministic fixes for known error types
    2. Context-aware: heuristics based on error context (tool, action, timing)
    3. AI-powered: Claude diagnosis for unknown or complex failures
    4. Default: generic retry as fallback
    """

    # ── Rule-based fix table ───────────────────────────────────────

    RULE_FIXES: dict[str, dict[str, Any]] = {
        "timeout": {
            "fix_type": "retry",
            "fix_params": {"multiply_timeout": 2},
            "confidence": 0.7,
            "root_cause": "Operation exceeded time limit",
        },
        "rate_limit": {
            "fix_type": "retry",
            "fix_params": {"delay_ms": 5000, "backoff": "exponential"},
            "confidence": 0.8,
            "root_cause": "API rate limit exceeded",
        },
        "network": {
            "fix_type": "retry",
            "fix_params": {"delay_ms": 2000, "max_retries": 3},
            "confidence": 0.6,
            "root_cause": "Transient network failure",
        },
        "not_found": {
            "fix_type": "modify_params",
            "fix_params": {"check_url": True, "verify_path": True},
            "confidence": 0.5,
            "root_cause": "Resource not found at expected location",
        },
        "validation": {
            "fix_type": "modify_params",
            "fix_params": {"revalidate": True},
            "confidence": 0.6,
            "root_cause": "Input validation failed",
        },
        "auth": {
            "fix_type": "escalate",
            "fix_params": {"refresh_auth": True},
            "confidence": 0.3,
            "root_cause": "Authentication/authorization failure",
        },
        "permission": {
            "fix_type": "escalate",
            "fix_params": {},
            "confidence": 0.3,
            "root_cause": "Insufficient permissions",
        },
        "resource": {
            "fix_type": "retry",
            "fix_params": {"reduce_batch": True, "reduce_factor": 0.5},
            "confidence": 0.5,
            "root_cause": "Resource exhaustion (memory, disk, quota)",
        },
        "conflict": {
            "fix_type": "retry",
            "fix_params": {"delay_ms": 1000, "randomize": True},
            "confidence": 0.6,
            "root_cause": "Resource conflict or deadlock",
        },
        "server_error": {
            "fix_type": "retry",
            "fix_params": {"delay_ms": 3000, "max_retries": 2},
            "confidence": 0.5,
            "root_cause": "Remote server error",
        },
    }

    # ── Context-aware heuristics ───────────────────────────────────

    TOOL_SPECIFIC_FIXES: dict[str, dict[str, dict[str, Any]]] = {
        "database": {
            "timeout": {
                "fix_type": "retry",
                "fix_params": {"multiply_timeout": 3, "reconnect": True},
                "confidence": 0.75,
                "root_cause": "Database query timeout — possible deadlock or slow query",
            },
            "network": {
                "fix_type": "retry",
                "fix_params": {"delay_ms": 1000, "reconnect": True},
                "confidence": 0.8,
                "root_cause": "Database connection lost",
            },
        },
        "http": {
            "timeout": {
                "fix_type": "retry",
                "fix_params": {"multiply_timeout": 2, "reduce_payload": True},
                "confidence": 0.7,
                "root_cause": "HTTP request timeout — server may be overloaded",
            },
            "validation": {
                "fix_type": "modify_params",
                "fix_params": {"accept": "application/json", "reformat": True},
                "confidence": 0.65,
                "root_cause": "HTTP response parse error",
            },
        },
        "browser": {
            "timeout": {
                "fix_type": "retry",
                "fix_params": {"restart_browser": True, "multiply_timeout": 2},
                "confidence": 0.7,
                "root_cause": "Browser operation timeout",
            },
            "unknown": {
                "fix_type": "retry",
                "fix_params": {"restart_browser": True},
                "confidence": 0.65,
                "root_cause": "Browser crashed or became unresponsive",
            },
        },
        "ai": {
            "validation": {
                "fix_type": "retry",
                "fix_params": {"stricter_prompt": True, "temperature": 0.0},
                "confidence": 0.6,
                "root_cause": "AI output did not meet validation requirements",
            },
            "resource": {
                "fix_type": "modify_params",
                "fix_params": {"reduce_context": True, "max_tokens": 2048},
                "confidence": 0.7,
                "root_cause": "AI context window exceeded",
            },
        },
        "email": {
            "network": {
                "fix_type": "retry",
                "fix_params": {"delay_ms": 5000},
                "confidence": 0.7,
                "root_cause": "SMTP connection failure",
            },
        },
        "files": {
            "resource": {
                "fix_type": "retry",
                "fix_params": {"delay_ms": 1000},
                "confidence": 0.75,
                "root_cause": "File lock or disk resource issue",
            },
            "permission": {
                "fix_type": "escalate",
                "fix_params": {"check_permissions": True},
                "confidence": 0.4,
                "root_cause": "Insufficient file permissions",
            },
        },
    }

    def __init__(
        self,
        ai_client: Any = None,
        model: str = "claude-sonnet-4-20250514",
    ) -> None:
        self.ai_client = ai_client
        self.model = model
        self._diagnosis_count: int = 0
        self._ai_fallback_count: int = 0
        self._cache: dict[str, Diagnosis] = {}
        self._max_cache: int = 200

    # ── Main diagnosis entry point ─────────────────────────────────

    async def diagnose(self, failure: Failure) -> dict[str, Any]:
        """Diagnose a failure and return fix recommendation.

        Priority:
        1. Check cache for recent identical failure
        2. Try tool-specific heuristics
        3. Try generic rule-based fixes
        4. Fall back to AI diagnosis
        5. Return default retry

        Returns a dict with: diagnosis, fix_type, fix_params, confidence
        """
        self._diagnosis_count += 1

        # Check cache first
        cache_key = f"{failure.error_type}:{failure.tool}:{failure.action}"
        cached = self._cache.get(cache_key)
        if cached and cached.confidence >= 0.5:
            return self._diagnosis_to_dict(cached)

        # Try tool-specific fix first (more precise)
        tool_fix = self._get_tool_specific_fix(failure)
        if tool_fix:
            diag = self._build_diagnosis(failure, tool_fix, source="rule")
            self._cache_diagnosis(cache_key, diag)
            return self._diagnosis_to_dict(diag)

        # Try generic rule-based fix
        rule_fix = self.RULE_FIXES.get(failure.error_type)
        if rule_fix:
            diag = self._build_diagnosis(failure, rule_fix, source="rule")
            self._cache_diagnosis(cache_key, diag)
            return self._diagnosis_to_dict(diag)

        # AI fallback
        if self.ai_client:
            self._ai_fallback_count += 1
            result = await self._ai_diagnose(failure)
            diag = self._build_diagnosis(failure, result, source="ai")
            self._cache_diagnosis(cache_key, diag)
            return self._diagnosis_to_dict(diag)

        # Default fallback
        default = {
            "diagnosis": f"Unknown error: {failure.error_message[:100]}",
            "fix_type": "retry",
            "fix_params": {},
            "confidence": 0.3,
            "root_cause": "Unclassified error",
        }
        return self._diagnosis_to_dict(
            self._build_diagnosis(failure, default, source="rule")
        )

    # ── Full diagnosis (returns Diagnosis object) ──────────────────

    async def diagnose_full(self, failure: Failure) -> Diagnosis:
        """Return a full Diagnosis object instead of a dict."""
        result = await self.diagnose(failure)
        return Diagnosis(
            failure=failure,
            root_cause=result.get("diagnosis", ""),
            fix_type=result.get("fix_type", "retry"),
            fix_params=result.get("fix_params", {}),
            confidence=result.get("confidence", 0.3),
            reasoning=result.get("root_cause", ""),
            source=result.get("source", "rule"),
        )

    # ── Multi-strategy diagnosis ───────────────────────────────────

    async def diagnose_multi_strategy(
        self, failure: Failure,
    ) -> list[dict[str, Any]]:
        """Generate multiple possible diagnoses ranked by confidence.

        Returns a list of fix options the caller can choose from.
        """
        strategies: list[dict[str, Any]] = []

        # Tool-specific fix
        tool_fix = self._get_tool_specific_fix(failure)
        if tool_fix:
            strategies.append({
                **tool_fix,
                "source": "tool_specific",
                "diagnosis": tool_fix.get("root_cause", "Tool-specific fix"),
            })

        # Generic rule fix
        rule_fix = self.RULE_FIXES.get(failure.error_type)
        if rule_fix:
            strategies.append({
                **rule_fix,
                "source": "rule",
                "diagnosis": rule_fix.get("root_cause", "Rule-based fix"),
            })

        # Skip option (always available for non-critical)
        if failure.recoverable:
            strategies.append({
                "fix_type": "skip",
                "fix_params": {},
                "confidence": 0.2,
                "source": "rule",
                "diagnosis": "Skip this step and continue",
                "root_cause": "Step can be safely skipped",
            })

        # Sort by confidence descending
        strategies.sort(key=lambda s: s.get("confidence", 0), reverse=True)
        return strategies

    # ── Tool-specific lookup ───────────────────────────────────────

    def _get_tool_specific_fix(self, failure: Failure) -> dict[str, Any] | None:
        """Look up a tool-specific fix for the failure."""
        tool_fixes = self.TOOL_SPECIFIC_FIXES.get(failure.tool)
        if not tool_fixes:
            return None
        return tool_fixes.get(failure.error_type)

    # ── AI diagnosis ───────────────────────────────────────────────

    async def _ai_diagnose(self, failure: Failure) -> dict[str, Any]:
        """Use AI to diagnose an unknown failure type."""
        prompt = f"""Diagnose this execution failure and suggest a fix.

Error type: {failure.error_type}
Error message: {failure.error_message}
Tool: {failure.tool}
Action: {failure.action}
Severity: {failure.severity}
Recoverable: {failure.recoverable}
Context: {json.dumps(failure.context, default=str)[:500]}

Respond with JSON:
{{"diagnosis": "root cause description",
  "fix_type": "retry|modify_params|skip|alternative_tool|escalate",
  "fix_params": {{}},
  "confidence": 0.0-1.0,
  "root_cause": "brief root cause"}}

Respond ONLY with valid JSON."""

        try:
            response = await self.ai_client.messages.create(
                model=self.model,
                max_tokens=1024,
                messages=[{"role": "user", "content": prompt}],
            )
            text = response.content[0].text.strip()
            # Strip markdown code fences if present
            for prefix in ["```json", "```"]:
                if text.startswith(prefix):
                    text = text[len(prefix):]
            if text.endswith("```"):
                text = text[:-3]
            result = json.loads(text.strip())
            # Validate fix_type
            valid_types = {"retry", "modify_params", "skip", "alternative_tool", "escalate"}
            if result.get("fix_type") not in valid_types:
                result["fix_type"] = "retry"
            return dict(result)
        except Exception:
            return {
                "diagnosis": "AI diagnosis failed — using default",
                "fix_type": "retry",
                "fix_params": {},
                "confidence": 0.3,
                "root_cause": "AI diagnosis unavailable",
            }

    # ── Internal helpers ───────────────────────────────────────────

    def _build_diagnosis(
        self, failure: Failure, fix_data: dict[str, Any], source: str,
    ) -> Diagnosis:
        """Build a Diagnosis from failure + fix data dict."""
        return Diagnosis(
            failure=failure,
            root_cause=fix_data.get("root_cause", fix_data.get("diagnosis", "")),
            fix_type=fix_data.get("fix_type", "retry"),
            fix_params=fix_data.get("fix_params", {}),
            confidence=fix_data.get("confidence", 0.5),
            reasoning=fix_data.get("diagnosis", ""),
            source=source,
        )

    def _diagnosis_to_dict(self, diag: Diagnosis) -> dict[str, Any]:
        """Convert a Diagnosis to the dict format expected by callers."""
        return {
            "diagnosis": diag.reasoning or diag.root_cause,
            "fix_type": diag.fix_type,
            "fix_params": diag.fix_params,
            "confidence": diag.confidence,
            "root_cause": diag.root_cause,
            "source": diag.source,
        }

    def _cache_diagnosis(self, key: str, diag: Diagnosis) -> None:
        """Cache a diagnosis for reuse on identical failures."""
        if len(self._cache) >= self._max_cache:
            # Evict oldest entries (first 25%)
            keys = list(self._cache.keys())
            for k in keys[:len(keys) // 4]:
                del self._cache[k]
        self._cache[key] = diag

    # ── Statistics ─────────────────────────────────────────────────

    def get_stats(self) -> dict[str, Any]:
        """Return diagnosis statistics."""
        return {
            "total_diagnoses": self._diagnosis_count,
            "ai_fallbacks": self._ai_fallback_count,
            "cache_size": len(self._cache),
            "ai_fallback_rate": (
                self._ai_fallback_count / self._diagnosis_count
                if self._diagnosis_count > 0
                else 0.0
            ),
        }

    def clear_cache(self) -> None:
        """Clear the diagnosis cache."""
        self._cache.clear()
