"""Vision AI adapter — screenshot analysis via multimodal models.

Supported providers:
  - aliyun  (qwen-vl-plus, qwen-vl-max)
  - 88api   (gpt-4o)

DeepSeek does NOT support vision in the public API.
API key is shared with the regular AI provider (loaded from keyring / ENV).
"""

from __future__ import annotations

import json
import logging
import os
import threading
from collections.abc import Iterator
from typing import TYPE_CHECKING

from ..security import privacy_manager

if TYPE_CHECKING:
    import openai as _openai

    from ..config import Config

log = logging.getLogger(__name__)

# Pattern to strip API keys and URL query parameters from error messages.
_SENSITIVE_RE = __import__("re").compile(
    r"((?:key|token|secret|password|api.?key)[=:\s]\s*)(\S+)"
    r"|(\?[^\s'\")\]]+)",
    __import__("re").IGNORECASE,
)


def _sanitize_error(exc: Exception) -> str:
    """Return a sanitized error string with API keys and query params removed."""
    msg = str(exc)
    msg = _SENSITIVE_RE.sub(
        lambda m: (m.group(1) + "***") if m.group(1) else "?<redacted>",
        msg,
    )
    return msg

# Vision-capable providers: name → (base_url, env_var, default_model)
VISION_PROVIDERS: dict[str, dict[str, str]] = {
    "aliyun": {
        "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1/",
        "env_var": "DASHSCOPE_API_KEY",
        "default_model": "qwen-vl-max",
    },
    "88api": {
        "base_url": "https://api.88api.dev/v1/",
        "env_var": "API88_API_KEY",
        "default_model": "gpt-4o",
    },
    # Local Ollama — uses config.ollama_base_url, no API key required.
    # Requires a vision-capable model: minicpm-v, llava, etc.
    "ollama": {
        "base_url": "http://localhost:11434/v1",  # overridden at runtime via config
        "env_var": "",
        "default_model": "qwen2.5-vl:7b",
    },
    # MuraveyChina license proxy — routes qwen-vl-* to Aliyun via server-side key.
    "proxy": {
        "base_url": "https://license.muravey.app/v1",
        "env_var": "",
        "default_model": "qwen-vl-max",
    },
    # llama-cpp-python — fully in-process, NO network connections at all.
    # Used when CSCA exam software blocks all ports (including localhost:11434).
    "llamacpp": {
        "base_url": "",  # no network — everything in-process
        "env_var": "",
        "default_model": "minicpm-v-2.6",
    },
}

_VISION_MODELS: dict[str, list[str]] = {
    "aliyun": ["qwen-vl-plus", "qwen-vl-max"],
    "88api": ["gpt-4o"],
    "ollama": ["qwen2.5-vl:7b", "minicpm-v", "llava", "llava:13b", "llava:7b"],
    "llamacpp": ["minicpm-v-2.6"],
}


class VisionError(Exception):
    """Raised when vision analysis cannot proceed."""


class VisionAdapter:
    """Sends screenshots to a Vision-capable AI provider and streams the response.

    Usage::

        adapter = VisionAdapter(config)
        for chunk in adapter.analyze(image_b64, "Describe this screen"):
            overlay.append(chunk)
    """

    def __init__(self, config: Config, timeout: float = 15.0) -> None:
        self._config = config
        self._timeout = timeout
        self._client_lock = threading.Lock()
        self._client: _openai.OpenAI | None = None
        self._last_provider: str = ""
        self._proxy_url: str = self._resolve_proxy_url()

    # ------------------------------------------------------------------ #
    # Public API
    # ------------------------------------------------------------------ #

    def analyze(
        self,
        image_b64: str,
        prompt: str,
        cancel_event: threading.Event | None = None,
        stream: bool = True,
        system_prompt: str | None = None,
        **kwargs,
    ) -> Iterator[str]:
        """Stream (or non-stream) AI analysis of a base64-encoded JPEG image.

        Args:
            image_b64: Base64-encoded JPEG string (no data URI prefix).
            prompt: Text prompt to send alongside the image.
            cancel_event: If set, streaming stops at the next chunk boundary.
            stream: If True (default), use SSE streaming. If False, use a single
                blocking request — faster for short responses (e.g. a single
                answer letter) because it avoids chunked-transfer overhead and
                saves ~50-150 ms of time-to-first-token RTT on China connections.
            **kwargs: Extra params for the API call (temperature, max_tokens, etc.).

        Yields:
            Text chunks as they arrive from the AI model.

        Raises:
            VisionError: If the provider is not supported or has no API key.
        """
        provider_name = self._config.vision_provider
        model = kwargs.pop("model", None) or self._config.vision_model

        if provider_name not in VISION_PROVIDERS:
            raise VisionError(
                f"Provider '{provider_name}' does not support vision. "
                f"Use: {', '.join(VISION_PROVIDERS)}"
            )

        api_key = self._resolve_key(provider_name)
        if not api_key:
            raise VisionError(
                f"No API key for provider '{provider_name}'. Set it in Settings → AI Provider."
            )

        log.info(
            "[vision] Sending screenshot to %s/%s (%d chars prompt, stream=%s)",
            provider_name,
            model,
            len(prompt),
            stream,
        )

        max_tokens = kwargs.get("max_tokens", 512)
        system_prompt_resolved = system_prompt if system_prompt is not None else (
            "You are an expert exam solver. You MUST reply with ONLY one single letter: A, B, C, or D. "
            "Do NOT write anything else. Do NOT explain. The ONLY valid responses are: A, B, C, or D."
        )

        # llama-cpp-python: fully in-process, no network at all (DEFAULT path)
        if provider_name == "llamacpp":
            yield from self._analyze_llamacpp(
                image_b64, prompt,
                system_prompt=system_prompt_resolved,
                max_tokens=max_tokens,
            )
            return

        # Ollama: use native /api/chat with images[] — fallback only
        if provider_name == "ollama":
            yield from self._analyze_ollama_native(
                image_b64, prompt,
                model=model,
                system_prompt=system_prompt_resolved,
                max_tokens=max_tokens,
                cancel_event=cancel_event,
            )
            return

        client = self._get_client(provider_name, api_key)
        messages = self._build_messages(image_b64, prompt, system_prompt=system_prompt)

        temperature = kwargs.get("temperature")

        create_kwargs: dict = {
            "model": model,
            "messages": messages,
            "stream": stream,
        }
        if temperature is not None:
            create_kwargs["temperature"] = temperature
        if max_tokens is not None:
            create_kwargs["max_tokens"] = max_tokens

        last_exc: Exception | None = None
        max_attempts = 2 if self._proxy_url else 1
        for attempt in range(max_attempts):
            try:
                if attempt > 0:
                    # Second attempt: drop proxy and connect directly
                    log.info("[vision] Retrying WITHOUT proxy (direct connection)...")
                    self._proxy_url = ""
                    self._client = None
                    client = self._get_client(provider_name, api_key)

                response = client.chat.completions.create(**create_kwargs)

                if stream:
                    for chunk in response:
                        if cancel_event and cancel_event.is_set():
                            log.info("[vision] Stream cancelled by user")
                            return
                        if chunk.choices and chunk.choices[0].delta.content:
                            yield chunk.choices[0].delta.content
                else:
                    text = response.choices[0].message.content or ""
                    if text:
                        yield text

                return
            except (ConnectionError, TimeoutError, OSError) as exc:
                last_exc = exc
                log.warning("[vision] Connection error (attempt %d/%d): %s",
                            attempt + 1, max_attempts, exc)
                continue
            except Exception as exc:
                exc_str = str(type(exc).__name__).lower()
                if "connection" in exc_str or "timeout" in exc_str or "Connection error" in str(exc):
                    last_exc = exc
                    log.warning("[vision] API error (attempt %d/%d): %s",
                                attempt + 1, max_attempts, exc)
                    continue
                log.error("[vision] Error: %s", exc)
                raise

        if last_exc:
            log.error("[vision] All attempts failed: %s", last_exc)
            raise VisionError(f"All attempts failed: {_sanitize_error(last_exc)}") from last_exc

    # ------------------------------------------------------------------ #
    # Ollama native API (more reliable for vision models than /v1)
    # ------------------------------------------------------------------ #

    def _analyze_ollama_native(
        self,
        image_b64: str,
        prompt: str,
        *,
        model: str,
        system_prompt: str,
        max_tokens: int = 512,
        cancel_event: threading.Event | None = None,
    ) -> Iterator[str]:
        """Call Ollama /api/chat with images[] array — bypasses /v1 vision bugs."""
        import urllib.request as _urllib_request
        import time as _time

        base_url = getattr(self._config, "ollama_base_url", "http://localhost:11434/v1")
        root = base_url.rstrip("/")
        if root.endswith("/v1"):
            root = root[:-3]
        url = f"{root}/api/chat"

        messages: list[dict] = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({
            "role": "user",
            "content": prompt,
            "images": [image_b64],
        })

        payload = json.dumps({
            "model": model,
            "messages": messages,
            "stream": False,
            "options": {"num_predict": max_tokens, "temperature": 0},
        }).encode()

        for attempt in range(3):
            if cancel_event and cancel_event.is_set():
                return
            try:
                if attempt > 0:
                    _time.sleep(1.0 * attempt)
                req = _urllib_request.Request(
                    url, data=payload,
                    headers={"Content-Type": "application/json"},
                )
                with _urllib_request.urlopen(req, timeout=self._timeout) as resp:
                    data: dict = json.loads(resp.read())
                text = data.get("message", {}).get("content", "")
                if text:
                    yield text
                return
            except Exception as exc:
                log.warning("[vision] Ollama native attempt %d/3 failed: %s", attempt + 1, exc)
                if attempt == 2:
                    log.error("[vision] Ollama native all attempts failed: %s", exc)
                    raise VisionError(f"Ollama failed: {_sanitize_error(exc)}") from exc

    # ------------------------------------------------------------------ #
    # llama-cpp-python (fully in-process, no network)
    # ------------------------------------------------------------------ #

    def _analyze_llamacpp(
        self,
        image_b64: str,
        prompt: str,
        *,
        system_prompt: str,
        max_tokens: int = 512,
    ) -> Iterator[str]:
        """Call llama-cpp-python in-process — zero network connections."""
        try:
            from .provider_llamacpp import get_llamacpp_provider
        except ImportError:
            raise VisionError("llama-cpp-python not installed")

        provider = get_llamacpp_provider()
        if provider is None or not provider.vision_ready:
            raise VisionError("llamacpp vision model not loaded")

        try:
            text = provider.generate(
                prompt,
                image_base64=image_b64,
                system_prompt=system_prompt,
                max_tokens=max_tokens,
                temperature=0.0,
            )
            if text:
                yield text
            else:
                raise VisionError("llamacpp returned empty response")
        except VisionError:
            raise
        except Exception as exc:
            raise VisionError(f"llamacpp error: {exc}") from exc

    # ------------------------------------------------------------------ #
    # Helpers
    # ------------------------------------------------------------------ #

    def get_available_providers(self) -> list[str]:
        """Return list of vision providers that have valid API keys."""
        available = []
        for name in VISION_PROVIDERS:
            key = self._resolve_key(name)
            if key:
                available.append(name)
        return available

    @staticmethod
    def get_models_for_provider(provider: str) -> list[str]:
        """Return the list of vision-capable models for a provider."""
        return list(_VISION_MODELS.get(provider, []))

    @staticmethod
    def supported_providers() -> list[str]:
        """Return names of all vision-capable providers."""
        return list(VISION_PROVIDERS.keys())

    # ------------------------------------------------------------------ #
    # Internal
    # ------------------------------------------------------------------ #

    def _resolve_key(self, provider: str) -> str:
        """Load API key: keyring first, then ENV variable. Ollama/proxy/llamacpp need no user key."""
        if provider == "llamacpp":
            return "llamacpp"  # no key needed — fully in-process
        if provider == "ollama":
            return "ollama"
        if provider == "proxy":
            from ..license.license_manager_v2 import load_session
            session = load_session()
            return session.get("session_key", "") if session else ""
        key = privacy_manager.load_api_key(provider)
        if key:
            return key
        env_var = VISION_PROVIDERS[provider]["env_var"]
        return os.environ.get(env_var, "")

    def _resolve_proxy_url(self) -> str:
        """Resolve proxy URL from config + auto-detection + NetworkProtector."""
        manual = getattr(self._config, "proxy_url", "")
        if manual:
            return manual

        # Use NetworkProtector proxy only if verified (can actually forward traffic).
        # Prevents broken SOCKS5 from causing Connection errors.
        try:
            from ..utils.network_protect import _active_protector
            if (_active_protector and _active_protector.proxy_url
                    and getattr(_active_protector, "_proxy_alive", False)):
                return _active_protector.proxy_url
        except ImportError:
            pass

        if not getattr(self._config, "proxy_auto_detect", True):
            return ""
        try:
            from ..utils.gfw_detect import get_proxy_if_needed
            return get_proxy_if_needed()
        except Exception:
            return ""

    def _make_http_client(self) -> object | None:
        """Create httpx.Client with proxy + keepalive for vision API calls."""
        # Re-resolve proxy if it was empty (NetworkProtector may have started)
        if not self._proxy_url:
            self._proxy_url = self._resolve_proxy_url()
        if not self._proxy_url:
            return None
        try:
            import httpx
            import httpcore

            transport = httpcore.ConnectionPool(
                keepalive_expiry=60.0,
                retries=2,
            )
            return httpx.Client(
                proxy=self._proxy_url,
                timeout=httpx.Timeout(self._timeout, connect=15.0),
                transport=transport,
            )
        except ImportError:
            log.debug("[vision] httpx not installed — proxy will not be used")
            return None
        except Exception as exc:
            log.debug("[vision] Failed to create proxy http_client: %s", exc)
            return None

    def _get_client(self, provider: str, api_key: str):
        """Return (or create) an openai.OpenAI client for the given provider."""
        import openai

        with self._client_lock:
            # Re-resolve proxy — invalidate client if proxy appeared
            new_proxy = self._resolve_proxy_url()
            if new_proxy and new_proxy != self._proxy_url:
                log.info("[vision] Proxy now available: %s — recreating client", new_proxy)
                self._proxy_url = new_proxy
                self._client = None

            if self._client is None or self._last_provider != provider:
                if provider == "ollama":
                    base_url = self._config.ollama_base_url
                elif provider == "proxy":
                    base_url = getattr(self._config, "proxy_base_url", VISION_PROVIDERS["proxy"]["base_url"])
                else:
                    base_url = VISION_PROVIDERS[provider]["base_url"]

                kwargs: dict = {
                    "api_key": api_key,
                    "base_url": base_url,
                    "timeout": self._timeout,
                    "max_retries": 0,
                }
                # Add proxy http_client for online providers (not ollama)
                if provider != "ollama":
                    http_client = self._make_http_client()
                    if http_client:
                        kwargs["http_client"] = http_client

                self._client = openai.OpenAI(**kwargs)
                self._last_provider = provider
                log.debug("[vision] Client created for provider '%s' at %s (proxy=%s)",
                          provider, base_url, bool(self._proxy_url))
        return self._client

    @staticmethod
    def _build_messages(image_b64: str, prompt: str, system_prompt: str | None = None) -> list[dict]:
        """Build multimodal message list (OpenAI vision format)."""
        if system_prompt is None:
            system_prompt = "You are an expert exam solver. You MUST reply with ONLY one single letter: A, B, C, or D. Do NOT write anything else. Do NOT explain. The ONLY valid responses are: A, B, C, or D."
        return [
            {
                "role": "system",
                "content": system_prompt,
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{image_b64}",
                        },
                    },
                    {
                        "type": "text",
                        "text": prompt,
                    },
                ],
            }
        ]
