"""
AI provider abstraction — supports OpenAI-compatible APIs, Google Gemini,
and Anthropic, with automatic fallback across all configured providers.

Key improvements over competitors:
  ✓ All providers share ONE interface — swap/chain without config changes
  ✓ Structured JSON response parsing (not brittle regex)
  ✓ Confidence scores in every response
  ✓ Multiple ranked candidates — use the best one
  ✓ Vision support: providers without a vision model gracefully skip
  ✓ Rate-limit / transient error detection → triggers provider fallback
"""

from __future__ import annotations

import base64
import json
import logging
import re
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

import httpx

from playwright_healer.config import AIProvider, AIProviderConfig

logger = logging.getLogger("playwright_healer.ai")


# -----------------------------------------------------------------
# Response model
# -----------------------------------------------------------------

@dataclass
class HealingCandidate:
    """A single healed selector candidate returned by an AI provider."""

    selector: str
    selector_type: str      # css | xpath | role | text | placeholder | label | testid
    confidence: float       # 0.0 – 1.0
    reasoning: str = ""


@dataclass
class AIHealingResponse:
    """Parsed AI response."""

    candidates: List[HealingCandidate] = field(default_factory=list)
    raw_text: str = ""
    provider: str = ""
    model: str = ""
    latency_ms: float = 0.0
    tokens_used: int = 0

    @property
    def best(self) -> Optional[HealingCandidate]:
        if not self.candidates:
            return None
        return max(self.candidates, key=lambda c: c.confidence)


# -----------------------------------------------------------------
# Prompt helpers
# -----------------------------------------------------------------

_DOM_SYSTEM_PROMPT = """\
You are a Playwright test automation expert specialising in self-healing element locators.
Given a broken CSS/XPath selector, the page's DOM, and an element description, identify
the element and return the BEST Playwright locators for it.

RULES:
1. Return ONLY valid JSON — no markdown fences, no prose.
2. Prefer ARIA-based locators (role, text, placeholder) over CSS/XPath.
3. Return up to 3 ranked candidates, most confident first.
4. Confidence is a float 0.0–1.0.
5. selector_type must be one of: css | xpath | role | text | placeholder | label | testid
6. For "role" type, format selector as: ROLE::NAME  e.g. button::Submit
7. For "text" type, use the visible text content.

JSON schema:
{
  "candidates": [
    {
      "selector": "...",
      "selector_type": "css|xpath|role|text|placeholder|label|testid",
      "confidence": 0.95,
      "reasoning": "brief explanation"
    }
  ]
}
"""

_DOM_USER_TEMPLATE = """\
Broken selector: {selector}
Element description: {description}
Current page URL: {url}

Relevant DOM (truncated to 8000 chars):
{dom}
"""

_VISUAL_SYSTEM_PROMPT = """\
You are a Playwright test automation expert. Given a screenshot and a broken element
description, identify the element visually and suggest the best Playwright locators.

RULES:
1. Return ONLY valid JSON.
2. Prefer ARIA-based selectors (role, text, placeholder).
3. Up to 3 ranked candidates.
4. Confidence 0.0–1.0.

JSON schema:
{
  "candidates": [
    {
      "selector": "...",
      "selector_type": "css|xpath|role|text|placeholder|label|testid",
      "confidence": 0.85,
      "reasoning": "brief explanation"
    }
  ]
}
"""

_VISUAL_USER_TEMPLATE = """\
Broken selector: {selector}
Element description: {description}
Current page URL: {url}
"""


def _parse_ai_response(raw: str) -> List[HealingCandidate]:
    """Robustly parse JSON candidates from AI output."""
    # Strip markdown code fences if present
    clean = re.sub(r"```(?:json)?", "", raw).strip().rstrip("```").strip()
    try:
        data = json.loads(clean)
        candidates = []
        for c in data.get("candidates", []):
            candidates.append(
                HealingCandidate(
                    selector=c.get("selector", "").strip(),
                    selector_type=c.get("selector_type", "css"),
                    confidence=float(c.get("confidence", 0.5)),
                    reasoning=c.get("reasoning", ""),
                )
            )
        return candidates
    except (json.JSONDecodeError, TypeError, KeyError) as exc:
        logger.debug("Could not parse AI response as JSON: %s — raw: %s", exc, raw[:200])
        return []


# -----------------------------------------------------------------
# Base provider
# -----------------------------------------------------------------

class AbstractAIProvider(ABC):
    """Base class for AI providers."""

    def __init__(self, cfg: AIProviderConfig) -> None:
        self._cfg = cfg

    @property
    def name(self) -> str:
        return self._cfg.provider.value

    @abstractmethod
    async def heal_with_dom(
        self,
        selector: str,
        description: str,
        dom: str,
        url: str,
    ) -> AIHealingResponse:
        ...

    @abstractmethod
    async def heal_with_vision(
        self,
        selector: str,
        description: str,
        screenshot_b64: str,
        url: str,
    ) -> AIHealingResponse:
        ...

    def supports_vision(self) -> bool:
        return bool(self._cfg.vision_model)

    @staticmethod
    def _is_transient_error(status_code: int) -> bool:
        return status_code in (429, 500, 502, 503, 504)


# -----------------------------------------------------------------
# OpenAI-compatible provider (OpenAI, Groq, DeepSeek, Local)
# -----------------------------------------------------------------

class OpenAICompatibleProvider(AbstractAIProvider):
    """Handles all OpenAI-compatible chat completion endpoints."""

    async def _call(
        self,
        messages: List[Dict[str, Any]],
        model: str,
        timeout: float,
    ) -> tuple[str, int, float]:
        start = time.monotonic()
        headers = {
            "Authorization": f"Bearer {self._cfg.api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "model": model,
            "messages": messages,
            "temperature": self._cfg.temperature,
            "max_tokens": self._cfg.max_tokens,
            "response_format": {"type": "json_object"},
        }
        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.post(self._cfg.api_url, headers=headers, json=payload)  # type: ignore[arg-type]
        latency = (time.monotonic() - start) * 1000
        if resp.status_code != 200:
            if self._is_transient_error(resp.status_code):
                raise TransientAIError(f"{self.name} returned {resp.status_code}")
            resp.raise_for_status()
        body = resp.json()
        content: str = body["choices"][0]["message"]["content"]
        tokens: int = body.get("usage", {}).get("total_tokens", 0)
        return content, tokens, latency

    async def heal_with_dom(
        self,
        selector: str,
        description: str,
        dom: str,
        url: str,
    ) -> AIHealingResponse:
        msgs = [
            {"role": "system", "content": _DOM_SYSTEM_PROMPT},
            {
                "role": "user",
                "content": _DOM_USER_TEMPLATE.format(
                    selector=selector,
                    description=description,
                    url=url,
                    dom=dom[:8000],
                ),
            },
        ]
        raw, tokens, latency = await self._call(msgs, self._cfg.model, self._cfg.timeout_s)
        return AIHealingResponse(
            candidates=_parse_ai_response(raw),
            raw_text=raw,
            provider=self.name,
            model=self._cfg.model,
            latency_ms=latency,
            tokens_used=tokens,
        )

    async def heal_with_vision(
        self,
        selector: str,
        description: str,
        screenshot_b64: str,
        url: str,
    ) -> AIHealingResponse:
        if not self._cfg.vision_model:
            return AIHealingResponse(provider=self.name, model=self._cfg.model)
        msgs = [
            {"role": "system", "content": _VISUAL_SYSTEM_PROMPT},
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": _VISUAL_USER_TEMPLATE.format(
                        selector=selector, description=description, url=url
                    )},
                    {"type": "image_url", "image_url": {
                        "url": f"data:image/png;base64,{screenshot_b64}"
                    }},
                ],
            },
        ]
        raw, tokens, latency = await self._call(msgs, self._cfg.vision_model, self._cfg.timeout_s)
        return AIHealingResponse(
            candidates=_parse_ai_response(raw),
            raw_text=raw,
            provider=self.name,
            model=self._cfg.vision_model,
            latency_ms=latency,
            tokens_used=tokens,
        )


# -----------------------------------------------------------------
# Gemini provider
# -----------------------------------------------------------------

class GeminiProvider(AbstractAIProvider):
    """Google Gemini via REST API (no SDK dependency)."""

    def _endpoint(self, model: str) -> str:
        base = "https://generativelanguage.googleapis.com/v1/models"
        return f"{base}/{model}:generateContent?key={self._cfg.api_key}"

    async def _call(
        self,
        parts: List[Dict[str, Any]],
        model: str,
        timeout: float,
    ) -> tuple[str, int, float]:
        start = time.monotonic()
        payload = {
            "contents": [{"role": "user", "parts": parts}],
            "systemInstruction": {"parts": [{"text": _DOM_SYSTEM_PROMPT}]},
            "generationConfig": {
                "temperature": self._cfg.temperature,
                "maxOutputTokens": self._cfg.max_tokens,
                "responseMimeType": "application/json",
            },
        }
        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.post(self._endpoint(model), json=payload)
        latency = (time.monotonic() - start) * 1000
        if resp.status_code != 200:
            if self._is_transient_error(resp.status_code):
                raise TransientAIError(f"Gemini returned {resp.status_code}")
            resp.raise_for_status()
        body = resp.json()
        content = body["candidates"][0]["content"]["parts"][0]["text"]
        tokens = body.get("usageMetadata", {}).get("totalTokenCount", 0)
        return content, tokens, latency

    async def heal_with_dom(
        self,
        selector: str,
        description: str,
        dom: str,
        url: str,
    ) -> AIHealingResponse:
        parts = [{"text": _DOM_USER_TEMPLATE.format(
            selector=selector, description=description, url=url, dom=dom[:8000]
        )}]
        raw, tokens, latency = await self._call(parts, self._cfg.model, self._cfg.timeout_s)
        return AIHealingResponse(
            candidates=_parse_ai_response(raw),
            raw_text=raw,
            provider=self.name,
            model=self._cfg.model,
            latency_ms=latency,
            tokens_used=tokens,
        )

    async def heal_with_vision(
        self,
        selector: str,
        description: str,
        screenshot_b64: str,
        url: str,
    ) -> AIHealingResponse:
        parts = [
            {"text": _VISUAL_USER_TEMPLATE.format(
                selector=selector, description=description, url=url
            )},
            {"inlineData": {"mimeType": "image/png", "data": screenshot_b64}},
        ]
        model = self._cfg.vision_model or self._cfg.model
        raw, tokens, latency = await self._call(parts, model, self._cfg.timeout_s)
        return AIHealingResponse(
            candidates=_parse_ai_response(raw),
            raw_text=raw,
            provider=self.name,
            model=model,
            latency_ms=latency,
            tokens_used=tokens,
        )


# -----------------------------------------------------------------
# Anthropic provider
# -----------------------------------------------------------------

class AnthropicProvider(AbstractAIProvider):
    """Anthropic Claude via Messages API."""

    async def _call(
        self,
        content: Any,
        model: str,
        system: str,
        timeout: float,
    ) -> tuple[str, int, float]:
        start = time.monotonic()
        headers = {
            "x-api-key": self._cfg.api_key,
            "anthropic-version": "2023-06-01",
            "content-type": "application/json",
        }
        payload = {
            "model": model,
            "max_tokens": self._cfg.max_tokens,
            "system": system,
            "messages": [{"role": "user", "content": content}],
        }
        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.post(
                "https://api.anthropic.com/v1/messages",
                headers=headers,
                json=payload,
            )
        latency = (time.monotonic() - start) * 1000
        if resp.status_code != 200:
            if self._is_transient_error(resp.status_code):
                raise TransientAIError(f"Anthropic returned {resp.status_code}")
            resp.raise_for_status()
        body = resp.json()
        text: str = body["content"][0]["text"]
        tokens: int = body.get("usage", {}).get("input_tokens", 0) + body.get("usage", {}).get("output_tokens", 0)
        return text, tokens, latency

    async def heal_with_dom(
        self,
        selector: str,
        description: str,
        dom: str,
        url: str,
    ) -> AIHealingResponse:
        user_text = _DOM_USER_TEMPLATE.format(
            selector=selector, description=description, url=url, dom=dom[:8000]
        )
        raw, tokens, latency = await self._call(
            user_text, self._cfg.model, _DOM_SYSTEM_PROMPT, self._cfg.timeout_s
        )
        return AIHealingResponse(
            candidates=_parse_ai_response(raw),
            raw_text=raw,
            provider=self.name,
            model=self._cfg.model,
            latency_ms=latency,
            tokens_used=tokens,
        )

    async def heal_with_vision(
        self,
        selector: str,
        description: str,
        screenshot_b64: str,
        url: str,
    ) -> AIHealingResponse:
        model = self._cfg.vision_model or self._cfg.model
        content = [
            {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/png",
                    "data": screenshot_b64,
                },
            },
            {"type": "text", "text": _VISUAL_USER_TEMPLATE.format(
                selector=selector, description=description, url=url
            )},
        ]
        raw, tokens, latency = await self._call(
            content, model, _VISUAL_SYSTEM_PROMPT, self._cfg.timeout_s
        )
        return AIHealingResponse(
            candidates=_parse_ai_response(raw),
            raw_text=raw,
            provider=self.name,
            model=model,
            latency_ms=latency,
            tokens_used=tokens,
        )


# -----------------------------------------------------------------
# Errors
# -----------------------------------------------------------------

class TransientAIError(Exception):
    """Raised for rate limits and server errors — triggers provider fallback."""


# -----------------------------------------------------------------
# Provider factory + fallback chain
# -----------------------------------------------------------------

def _build_provider(cfg: AIProviderConfig) -> AbstractAIProvider:
    mapping = {
        AIProvider.OPENAI: OpenAICompatibleProvider,
        AIProvider.GROQ: OpenAICompatibleProvider,
        AIProvider.DEEPSEEK: OpenAICompatibleProvider,
        AIProvider.LOCAL: OpenAICompatibleProvider,
        AIProvider.GEMINI: GeminiProvider,
        AIProvider.ANTHROPIC: AnthropicProvider,
    }
    cls = mapping.get(cfg.provider, OpenAICompatibleProvider)
    return cls(cfg)


class ProviderChain:
    """
    Wraps multiple providers; on TransientAIError automatically falls back
    to the next one in the list.  This is a key advantage: competitor
    packages crash when a single provider hits a rate limit.
    """

    def __init__(self, configs: List[AIProviderConfig]) -> None:
        self._providers = [_build_provider(c) for c in configs]
        if not self._providers:
            raise ValueError("ProviderChain requires at least one AIProviderConfig")

    async def heal_with_dom(
        self,
        selector: str,
        description: str,
        dom: str,
        url: str,
    ) -> AIHealingResponse:
        last_exc: Optional[Exception] = None
        for provider in self._providers:
            try:
                result = await provider.heal_with_dom(selector, description, dom, url)
                if result.candidates:
                    return result
            except (TransientAIError, httpx.HTTPError, Exception) as exc:
                logger.warning(
                    "Provider %s failed (DOM healing): %s — trying next provider",
                    provider.name, exc
                )
                last_exc = exc
        raise RuntimeError(
            f"All AI providers exhausted for selector '{selector}'"
        ) from last_exc

    async def heal_with_vision(
        self,
        selector: str,
        description: str,
        screenshot_b64: str,
        url: str,
    ) -> AIHealingResponse:
        last_exc: Optional[Exception] = None
        for provider in self._providers:
            if not provider.supports_vision():
                continue
            try:
                result = await provider.heal_with_vision(
                    selector, description, screenshot_b64, url
                )
                if result.candidates:
                    return result
            except (TransientAIError, httpx.HTTPError, Exception) as exc:
                logger.warning(
                    "Provider %s failed (visual healing): %s — trying next provider",
                    provider.name, exc
                )
                last_exc = exc
        raise RuntimeError(
            f"All vision AI providers exhausted for selector '{selector}'"
        ) from last_exc

    def has_vision(self) -> bool:
        return any(p.supports_vision() for p in self._providers)
