"""
Cloud Bridge — Calls Claude, Gemini, and GPT APIs with OAuth or BYOK credentials.

Credential resolution priority:
    1. OAuth token (from subscription login via OAuthManager)
    2. BYOK API key (from local vault / KeyVault)
    3. Error with helpful message

Makes direct HTTPS calls to provider APIs using the user's own credentials,
which are retrieved from local stores and NEVER sent to Brynq servers.

Backend-specific API formatting:
  - Claude: Anthropic Messages API (api.anthropic.com)
  - Gemini: GenerateContent API (generativelanguage.googleapis.com)
  - GPT: Chat Completions API (api.openai.com)

All stdlib — no external packages (except local project imports).
"""

from __future__ import annotations

import json
import logging
import socket
import ssl
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import Any, Generator, Optional

from .local_bridge import BridgeResponse, _estimate_tokens

log = logging.getLogger("brynq.runtime.cloud_bridge")

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages"
_ANTHROPIC_VERSION = "2023-06-01"

_GEMINI_API_BASE = "https://generativelanguage.googleapis.com/v1beta/models"

_OPENAI_API_URL = "https://api.openai.com/v1/chat/completions"

_DEFAULT_MAX_TOKENS = 2048
_HTTP_TIMEOUT = 120  # seconds
_RETRY_MAX_ATTEMPTS = 3
_RETRY_BASE_DELAY = 1  # seconds; exponential backoff: 1s, 4s, 16s


# ---------------------------------------------------------------------------
# Credential types
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class Credential:
    """A resolved credential for a cloud provider API call.

    Attributes:
        type: Either "oauth" or "api_key".
        value: The token or key value. NEVER logged.
        provider: Normalized provider name (e.g. "anthropic", "google", "openai").
    """

    type: str  # "oauth" | "api_key"
    value: str
    provider: str


class NoCredentialsError(RuntimeError):
    """Raised when no OAuth token or API key is available for a provider."""


# ---------------------------------------------------------------------------
# Provider name mapping
# ---------------------------------------------------------------------------

# Maps backend names used in calls to the provider names used by
# OAuthManager and KeyVault.
_BACKEND_TO_PROVIDER: dict[str, str] = {
    "claude": "anthropic",
    "gemini": "google",
    "gpt": "openai",
    "openai": "openai",
}


# ---------------------------------------------------------------------------
# Internal HTTP helpers
# ---------------------------------------------------------------------------


def _build_ssl_context() -> ssl.SSLContext:
    """Return a default SSL context for HTTPS requests."""
    return ssl.create_default_context()


def _http_json(
    method: str,
    url: str,
    body: Optional[dict] = None,
    headers: Optional[dict[str, str]] = None,
    timeout: int = _HTTP_TIMEOUT,
) -> dict:
    """Make an HTTP request and return parsed JSON.

    Raises RuntimeError on HTTP or network errors with descriptive messages.
    """
    hdrs: dict[str, str] = {
        "Content-Type": "application/json",
        "Accept": "application/json",
    }
    if headers:
        hdrs.update(headers)

    data = json.dumps(body).encode("utf-8") if body is not None else None
    req = urllib.request.Request(url, data=data, headers=hdrs, method=method)
    context = _build_ssl_context() if url.startswith("https://") else None

    try:
        with urllib.request.urlopen(req, timeout=timeout, context=context) as resp:
            raw = resp.read().decode("utf-8")
            return json.loads(raw) if raw.strip() else {}
    except urllib.error.HTTPError as exc:
        detail = ""
        try:
            detail = exc.read().decode("utf-8", errors="replace")
        except Exception:
            pass
        raise RuntimeError(f"HTTP {exc.code} {method} {url}: {detail}") from exc
    except urllib.error.URLError as exc:
        if isinstance(exc.reason, socket.gaierror):
            # Extract provider name from URL for a user-friendly message
            host = url.split("//", 1)[-1].split("/", 1)[0].split(":")[0]
            raise RuntimeError(
                f"Cannot resolve {host} API hostname. Check your internet connection."
            ) from exc
        raise RuntimeError(f"Connection failed {method} {url}: {exc.reason}") from exc


def _http_json_with_retry(
    method: str,
    url: str,
    body: Optional[dict] = None,
    headers: Optional[dict[str, str]] = None,
    timeout: int = _HTTP_TIMEOUT,
) -> dict:
    """Wrap _http_json with retry on HTTP 429 (rate limit) responses.

    Retries up to 3 times with exponential backoff (1s, 4s, 16s delays).
    Logs each retry at WARNING level.
    """
    for attempt in range(_RETRY_MAX_ATTEMPTS + 1):
        try:
            return _http_json(method, url, body=body, headers=headers, timeout=timeout)
        except RuntimeError as exc:
            if "HTTP 429" in str(exc) and attempt < _RETRY_MAX_ATTEMPTS:
                delay = _RETRY_BASE_DELAY * (4 ** attempt)
                log.warning(
                    "Rate limited (429), retry %d/%d after %.1fs: %s %s",
                    attempt + 1, _RETRY_MAX_ATTEMPTS, delay, method, url,
                )
                time.sleep(delay)
            else:
                raise
    # Unreachable, but satisfies type checker
    raise RuntimeError("Retry loop exhausted")


# ---------------------------------------------------------------------------
# Backend-specific callers (support both OAuth and API key auth)
# ---------------------------------------------------------------------------


def _call_claude(
    model: str,
    prompt: str,
    credential: Credential,
    config: dict,
) -> BridgeResponse:
    """Call the Anthropic Messages API.

    Auth differences:
    - OAuth token:  ``Authorization: Bearer {token}``
    - API key:      ``x-api-key: {key}``
    """
    response = BridgeResponse(model=model, backend="claude")
    t0 = time.monotonic()

    try:
        if not credential.value:
            raise RuntimeError("Claude credential not provided")

        # Build messages
        messages: list[dict[str, str]] = []
        system_text: Optional[str] = config.get("system_prompt")

        messages.append({"role": "user", "content": prompt})

        body: dict[str, Any] = {
            "model": model,
            "max_tokens": config.get("max_tokens", _DEFAULT_MAX_TOKENS),
            "messages": messages,
        }
        if system_text:
            body["system"] = system_text
        if "temperature" in config:
            body["temperature"] = config["temperature"]

        # Auth header depends on credential type
        headers: dict[str, str] = {
            "anthropic-version": _ANTHROPIC_VERSION,
        }
        if credential.type == "oauth":
            headers["Authorization"] = f"Bearer {credential.value}"
        else:
            headers["x-api-key"] = credential.value

        resp = _http_json_with_retry(
            "POST",
            _ANTHROPIC_API_URL,
            body=body,
            headers=headers,
        )

        # Extract text from response blocks
        parts: list[str] = []
        for block in resp.get("content", []):
            if block.get("type") == "text":
                parts.append(block.get("text", ""))
        output = "\n".join(parts)

        response.output = output
        response.success = True
        response.token_estimate = _estimate_tokens(output)
        response.metadata["auth_type"] = credential.type

        # Log stop reason
        stop_reason = resp.get("stop_reason")
        if stop_reason:
            response.metadata["stop_reason"] = stop_reason
            log.debug("Claude stop reason: %s (model=%s)", stop_reason, model)

        # Extract usage metadata
        usage = resp.get("usage", {})
        if usage:
            response.metadata["input_tokens"] = usage.get("input_tokens", 0)
            response.metadata["output_tokens"] = usage.get("output_tokens", 0)

    except RuntimeError as exc:
        response.success = False
        error_str = str(exc)
        if "401" in error_str:
            if credential.type == "oauth":
                response.error = "Claude OAuth token is invalid or expired. Try reconnecting in Settings."
            else:
                response.error = "Claude API key is invalid. Check your key in Settings."
        elif "429" in error_str:
            response.error = "Claude rate limit exceeded. Try again in a moment."
        elif "529" in error_str or "overloaded" in error_str.lower():
            response.error = "Claude API is overloaded. Try again later."
        else:
            response.error = f"Claude API error: {exc}"
        log.error("Cloud bridge (Claude) error: %s", exc)

    except Exception as exc:
        response.success = False
        response.error = f"Unexpected error calling Claude: {exc}"
        log.error("Cloud bridge (Claude) unexpected error: %s", exc, exc_info=True)

    finally:
        response.elapsed_seconds = time.monotonic() - t0

    return response


def _call_gemini(
    model: str,
    prompt: str,
    credential: Credential,
    config: dict,
) -> BridgeResponse:
    """Call the Google Gemini GenerateContent API.

    Auth:
    - OAuth access token:  ``Authorization: Bearer {token}`` (subscription)
    - API key:             ``?key={key}`` query parameter (pay-per-token)

    OAuth tokens route through the user's Google subscription.
    API keys route to the Generative Language API billing project.
    """
    response = BridgeResponse(model=model, backend="gemini")
    t0 = time.monotonic()

    try:
        if not credential.value:
            raise RuntimeError("Gemini credential not provided")

        # Build contents
        contents: list[dict[str, Any]] = []
        system_prompt = config.get("system_prompt")

        contents.append({
            "role": "user",
            "parts": [{"text": prompt}],
        })

        body: dict[str, Any] = {"contents": contents}

        # System instruction (proper Gemini API field, not prepended to prompt)
        if system_prompt:
            body["systemInstruction"] = {
                "parts": [{"text": system_prompt}],
            }

        # Generation config
        gen_config: dict[str, Any] = {}
        if "temperature" in config:
            gen_config["temperature"] = config["temperature"]
        if "max_tokens" in config:
            gen_config["maxOutputTokens"] = config["max_tokens"]
        if gen_config:
            body["generationConfig"] = gen_config

        # Auth: OAuth uses Bearer header (subscription), API key uses ?key= (pay-per-token)
        base_url = f"{_GEMINI_API_BASE}/{model}:generateContent"
        headers: dict[str, str] = {}
        if credential.type == "oauth":
            url = base_url
            headers["Authorization"] = f"Bearer {credential.value}"
        else:
            url = f"{base_url}?key={credential.value}"

        resp = _http_json_with_retry("POST", url, body=body, headers=headers)

        # Extract text from candidates
        candidates = resp.get("candidates", [])
        text_parts: list[str] = []
        if candidates:
            for p in candidates[0].get("content", {}).get("parts", []):
                if "text" in p:
                    text_parts.append(p["text"])
        output = "\n".join(text_parts)

        response.output = output
        response.success = True
        response.token_estimate = _estimate_tokens(output)
        response.metadata["auth_type"] = credential.type

        # Log stop reason
        if candidates:
            finish_reason = candidates[0].get("finishReason")
            if finish_reason:
                response.metadata["stop_reason"] = finish_reason
                log.debug("Gemini stop reason: %s (model=%s)", finish_reason, model)

        # Extract usage metadata
        usage = resp.get("usageMetadata", {})
        if usage:
            response.metadata["prompt_tokens"] = usage.get("promptTokenCount", 0)
            response.metadata["output_tokens"] = usage.get("candidatesTokenCount", 0)

    except RuntimeError as exc:
        response.success = False
        error_str = str(exc)
        if "400" in error_str:
            response.error = "Gemini API request was invalid. Check model name and parameters."
        elif "403" in error_str:
            if credential.type == "oauth":
                response.error = "Gemini OAuth token is invalid or expired. Try reconnecting in Settings."
            else:
                response.error = "Gemini API key is invalid or lacks permissions."
        elif "429" in error_str:
            response.error = "Gemini rate limit exceeded. Try again in a moment."
        else:
            response.error = f"Gemini API error: {exc}"
        log.error("Cloud bridge (Gemini) error: %s", exc)

    except Exception as exc:
        response.success = False
        response.error = f"Unexpected error calling Gemini: {exc}"
        log.error("Cloud bridge (Gemini) unexpected error: %s", exc, exc_info=True)

    finally:
        response.elapsed_seconds = time.monotonic() - t0

    return response


def _call_gpt(
    model: str,
    prompt: str,
    credential: Credential,
    config: dict,
) -> BridgeResponse:
    """Call the OpenAI Chat Completions API.

    Auth note: Both OAuth and API key use the same
    ``Authorization: Bearer {value}`` header for OpenAI.
    """
    response = BridgeResponse(model=model, backend="gpt")
    t0 = time.monotonic()

    try:
        if not credential.value:
            raise RuntimeError("OpenAI credential not provided")

        # Build messages
        messages: list[dict[str, str]] = []
        system_prompt = config.get("system_prompt")
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        body: dict[str, Any] = {
            "model": model,
            "messages": messages,
        }
        if "temperature" in config:
            body["temperature"] = config["temperature"]
        if "max_tokens" in config:
            body["max_tokens"] = config["max_tokens"]

        # OpenAI: Both OAuth and API key use Bearer token
        resp = _http_json_with_retry(
            "POST",
            _OPENAI_API_URL,
            body=body,
            headers={
                "Authorization": f"Bearer {credential.value}",
            },
        )

        # Extract output
        choices = resp.get("choices", [])
        output = ""
        if choices:
            output = choices[0].get("message", {}).get("content", "")

        response.output = output
        response.success = True
        response.token_estimate = _estimate_tokens(output)
        response.metadata["auth_type"] = credential.type

        # Log stop reason
        if choices:
            finish_reason = choices[0].get("finish_reason")
            if finish_reason:
                response.metadata["stop_reason"] = finish_reason
                log.debug("OpenAI stop reason: %s (model=%s)", finish_reason, model)

        # Extract usage metadata
        usage = resp.get("usage", {})
        if usage:
            response.metadata["prompt_tokens"] = usage.get("prompt_tokens", 0)
            response.metadata["completion_tokens"] = usage.get("completion_tokens", 0)

    except RuntimeError as exc:
        response.success = False
        error_str = str(exc)
        if "401" in error_str:
            if credential.type == "oauth":
                response.error = "OpenAI OAuth token is invalid or expired. Try reconnecting in Settings."
            else:
                response.error = "OpenAI API key is invalid. Check your key in Settings."
        elif "429" in error_str:
            response.error = "OpenAI rate limit exceeded. Try again in a moment."
        elif "model_not_found" in error_str.lower() or "does not exist" in error_str.lower():
            response.error = f"Model '{model}' not found. Check the model name."
        else:
            response.error = f"OpenAI API error: {exc}"
        log.error("Cloud bridge (GPT) error: %s", exc)

    except Exception as exc:
        response.success = False
        response.error = f"Unexpected error calling OpenAI: {exc}"
        log.error("Cloud bridge (GPT) unexpected error: %s", exc, exc_info=True)

    finally:
        response.elapsed_seconds = time.monotonic() - t0

    return response


# ---------------------------------------------------------------------------
# Dispatch table
# ---------------------------------------------------------------------------

_BACKEND_DISPATCH = {
    "claude": _call_claude,
    "gemini": _call_gemini,
    "gpt": _call_gpt,
    "openai": _call_gpt,  # Alias
}


# ---------------------------------------------------------------------------
# Cloud Bridge
# ---------------------------------------------------------------------------


class CloudBridge:
    """Bridge to cloud LLM APIs (Claude, Gemini, GPT).

    Supports two credential sources with clear priority:
        1. OAuth token (from subscription login via OAuthManager)
        2. BYOK API key (from local KeyVault)
        3. Error with helpful message if neither available

    Credentials are read from local stores only and NEVER transmitted
    to Brynq servers.

    Usage with credential resolution (preferred)::

        from runtime.auth import OAuthManager
        from runtime.vault import KeyVault

        bridge = CloudBridge(oauth_manager=OAuthManager(), key_vault=KeyVault())
        result = await bridge.call_with_resolution("claude", "claude-sonnet-4-20250514", "Hello", {})

    Usage with explicit credential (backward compatible)::

        bridge = CloudBridge()
        result = bridge.call("claude", "claude-sonnet-4-20250514", "Hello", "sk-ant-...", {})
    """

    def __init__(
        self,
        oauth_manager: Optional[Any] = None,
        key_vault: Optional[Any] = None,
    ) -> None:
        """Initialize the CloudBridge.

        Args:
            oauth_manager: Optional OAuthManager instance for OAuth token resolution.
                           If None, OAuth credentials are not available.
            key_vault: Optional KeyVault instance for BYOK API key resolution.
                       If None, API key credentials are not available.
        """
        self._oauth_manager = oauth_manager
        self._key_vault = key_vault

    async def _get_credential(self, backend: str) -> Credential:
        """Resolve the best available credential for a backend.

        Priority:
            1. OAuth token (from subscription login)
            2. BYOK API key (from local vault)
            3. Error with helpful message

        Args:
            backend: Provider backend name ("claude", "gemini", "gpt", "openai").

        Returns:
            Credential with type, value, and provider.

        Raises:
            NoCredentialsError: If no credential source has a valid token/key.
        """
        provider = _BACKEND_TO_PROVIDER.get(backend, backend)

        # Priority 1: OAuth token
        if self._oauth_manager is not None:
            try:
                token = await self._oauth_manager.get_token(provider)
                if token:
                    log.debug("Using OAuth token for %s", provider)
                    return Credential(type="oauth", value=token, provider=provider)
            except Exception:
                log.debug("OAuth token retrieval failed for %s", provider, exc_info=True)

        # Priority 2: BYOK API key
        if self._key_vault is not None:
            try:
                api_key = self._key_vault.get_key(provider)
                if api_key:
                    log.debug("Using API key for %s", provider)
                    return Credential(type="api_key", value=api_key, provider=provider)
            except Exception:
                log.debug("API key retrieval failed for %s", provider, exc_info=True)

        # Priority 3: Error
        provider_display = {
            "anthropic": "Claude",
            "google": "Gemini",
            "openai": "ChatGPT",
        }.get(provider, provider)

        raise NoCredentialsError(
            f"No credentials available for {provider_display}. "
            f"Connect your {provider_display} account or add an API key in Settings."
        )

    async def call_with_resolution(
        self,
        backend: str,
        model: str,
        prompt: str,
        config: Optional[dict] = None,
    ) -> BridgeResponse:
        """Make an API call with automatic credential resolution.

        Resolves credentials using the priority chain (OAuth -> API key)
        then dispatches to the appropriate backend caller.

        Args:
            backend: Provider name -- "claude", "gemini", or "gpt"/"openai".
            model: Model identifier (e.g. "claude-sonnet-4-20250514").
            prompt: The user prompt to send.
            config: Optional config dict with keys like "temperature",
                    "max_tokens", "system_prompt".

        Returns:
            BridgeResponse with the model's output or error details.
        """
        config = config or {}

        handler = _BACKEND_DISPATCH.get(backend)
        if handler is None:
            return BridgeResponse(
                success=False,
                error=(
                    f"Unknown cloud backend '{backend}'. "
                    f"Supported: {', '.join(sorted(set(_BACKEND_DISPATCH.keys())))}"
                ),
                backend=backend,
                model=model,
            )

        try:
            credential = await self._get_credential(backend)
        except NoCredentialsError as exc:
            return BridgeResponse(
                success=False,
                error=str(exc),
                backend=backend,
                model=model,
            )

        result = handler(model, prompt, credential, config)

        # Retry once on 401 if we used an OAuth token: refresh and re-call
        if (
            not result.success
            and result.error
            and "401" in result.error
            and credential.type == "oauth"
            and self._oauth_manager is not None
        ):
            provider = _BACKEND_TO_PROVIDER.get(backend, backend)
            try:
                new_token = await self._oauth_manager.refresh_token(provider)
                log.info("Refreshed OAuth token for %s after 401, retrying", provider)
                credential = Credential(type="oauth", value=new_token, provider=provider)
                result = handler(model, prompt, credential, config)
            except Exception as refresh_exc:
                log.warning("OAuth token refresh failed for %s: %s", provider, refresh_exc)

        return result

    def call(
        self,
        backend: str,
        model: str,
        prompt: str,
        api_key: str,
        config: Optional[dict] = None,
    ) -> BridgeResponse:
        """Make a blocking API call with an explicit API key (backward compatible).

        This is the original BYOK interface. For automatic credential resolution
        (OAuth preferred over API key), use ``call_with_resolution()`` instead.

        Args:
            backend: Provider name -- "claude", "gemini", or "gpt"/"openai".
            model: Model identifier.
            prompt: The user prompt to send.
            api_key: The user's API key. NEVER sent to Brynq servers.
            config: Optional config dict.

        Returns:
            BridgeResponse with the model's output or error details.
        """
        config = config or {}

        handler = _BACKEND_DISPATCH.get(backend)
        if handler is None:
            return BridgeResponse(
                success=False,
                error=(
                    f"Unknown cloud backend '{backend}'. "
                    f"Supported: {', '.join(sorted(set(_BACKEND_DISPATCH.keys())))}"
                ),
                backend=backend,
                model=model,
            )

        provider = _BACKEND_TO_PROVIDER.get(backend, backend)
        credential = Credential(type="api_key", value=api_key, provider=provider)
        log.debug("Using API key for %s", provider)
        return handler(model, prompt, credential, config)

    def call_streaming(
        self,
        backend: str,
        model: str,
        prompt: str,
        api_key: str,
        config: Optional[dict] = None,
    ) -> Generator[str, None, None]:
        """Make a streaming API call to a cloud LLM provider.

        Currently falls back to a blocking call and yields the full output
        as a single chunk. True streaming (SSE) will be added per-backend
        in a future iteration.

        Args:
            backend: Provider name.
            model: Model identifier.
            prompt: The user prompt.
            api_key: The user's API key.
            config: Optional config dict.

        Yields:
            Text chunks from the model's response.
        """
        # For now, do a blocking call and yield the result.
        # True SSE streaming is a future enhancement.
        result = self.call(backend, model, prompt, api_key, config)
        if result.success:
            yield result.output
        else:
            yield f"[ERROR] {result.error}"
