"""Thin client for interacting with Google Gemini via Vertex AI."""

from __future__ import annotations

import concurrent.futures
import json
import logging
import os
import random
import re
import time
from typing import Any, Dict, List, Optional

import vertexai
from vertexai.generative_models import GenerativeModel, GenerationConfig

from .llm_guard import (
    LLMConcurrencyError,
    LLMCircuitOpenError,
    LLMGuard,
    record_failure,
    record_success,
)

# Import Google API exceptions for retry handling
try:
    from google.api_core import exceptions as google_exceptions
    HAS_GOOGLE_EXCEPTIONS = True
except ImportError:
    HAS_GOOGLE_EXCEPTIONS = False

logger = logging.getLogger(__name__)

# Retry configuration
MAX_RETRIES = 3
INITIAL_BACKOFF_S = 1.0
MAX_BACKOFF_S = 16.0
BACKOFF_MULTIPLIER = 2.0
_EXECUTOR = concurrent.futures.ThreadPoolExecutor(max_workers=8)

class GeminiClient:
    """Wrapper around Gemini via Vertex AI."""

    def __init__(
        self,
        model: str,
        timeout_s: int = 12,
        project_id: Optional[str] = None,
        location: str = "us-central1",
        max_output_tokens: Optional[int] = None,
        max_retries: Optional[int] = None,
        wall_clock_budget_s: Optional[int] = None,
        retry_jitter_min_ms: Optional[int] = None,
        retry_jitter_max_ms: Optional[int] = None,
        account_id: Optional[str] = None,
    ) -> None:
        """
        Initialize Gemini client using Vertex AI.

        Args:
            model: Gemini model name (e.g., "gemini-2.0-flash-exp")
            timeout_s: Request timeout in seconds
            project_id: GCP project ID (uses default from gcloud if None)
            location: GCP region
        """
        self.model_name = model
        self.timeout = timeout_s
        self.max_retries = (
            max_retries if max_retries is not None else max(0, MAX_RETRIES - 1)
        )
        self.wall_clock_budget_s = wall_clock_budget_s
        self.retry_jitter_min_ms = retry_jitter_min_ms
        self.retry_jitter_max_ms = retry_jitter_max_ms
        self.account_id = account_id
        if (model or "").strip().lower().startswith("gemini-3-"):
            location = "global"

        # Get project ID from env or gcloud config
        if project_id is None:
            project_id = os.getenv("GOOGLE_CLOUD_PROJECT") or "aqueous-encoder-478015-r7"

        self.project_id = project_id
        self.location = location

        # Initialize Vertex AI (uses gcloud auth)
        vertexai.init(project=project_id, location=location)

        if max_output_tokens is None:
            env_max = os.getenv("GEMINI_MAX_OUTPUT_TOKENS")
            if env_max:
                try:
                    max_output_tokens = int(env_max)
                except ValueError:
                    max_output_tokens = 2048
            else:
                max_output_tokens = 2048

        # Initialize model with JSON response config
        generation_config = GenerationConfig(
            temperature=0.1,  # Low for consistency
            top_p=0.9,
            top_k=40,
            max_output_tokens=max_output_tokens,
            response_mime_type="application/json",  # Force JSON output
        )

        self.model = GenerativeModel(
            model_name=model,
            generation_config=generation_config,
        )

    def _is_retryable_error(self, exc: Exception) -> bool:
        """Check if an exception is retryable (rate limit, service unavailable, etc.)."""
        if isinstance(exc, TimeoutError):
            return True
        if HAS_GOOGLE_EXCEPTIONS:
            # 429 Resource Exhausted (rate limit)
            if isinstance(exc, google_exceptions.ResourceExhausted):
                return True
            # 503 Service Unavailable
            if isinstance(exc, google_exceptions.ServiceUnavailable):
                return True
            # 500 Internal Server Error
            if isinstance(exc, google_exceptions.InternalServerError):
                return True
            # 504 Gateway Timeout
            if isinstance(exc, google_exceptions.DeadlineExceeded):
                return True

        # Fallback: check error message for common retry scenarios
        err_str = str(exc).lower()
        retryable_patterns = [
            "rate limit",
            "resource exhausted",
            "quota exceeded",
            "service unavailable",
            "503",
            "429",
            "timeout",
            "deadline exceeded",
        ]
        return any(pattern in err_str for pattern in retryable_patterns)

    def _sanitize_error(self, exc: Exception) -> str:
        """Sanitize error message to avoid exposing sensitive information."""
        err_str = str(exc)
        # Remove potential credential/token patterns
        err_str = re.sub(r"(Bearer\s+)[^\s]+", r"\1[REDACTED]", err_str)
        err_str = re.sub(r"(api[_-]?key[=:]\s*)[^\s&]+", r"\1[REDACTED]", err_str, flags=re.IGNORECASE)
        err_str = re.sub(r"(token[=:]\s*)[^\s&]+", r"\1[REDACTED]", err_str, flags=re.IGNORECASE)
        # Truncate long messages
        if len(err_str) > 500:
            err_str = err_str[:500] + "..."
        return err_str

    def _call_with_timeout(self, content: str, timeout_s: Optional[float]) -> Any:
        if timeout_s is None or timeout_s <= 0:
            return self.model.generate_content(content)
        future = _EXECUTOR.submit(self.model.generate_content, content)
        try:
            return future.result(timeout=timeout_s)
        except concurrent.futures.TimeoutError as exc:
            future.cancel()
            raise TimeoutError("Gemini request timed out") from exc

    def _retry_sleep_s(self, backoff: float) -> float:
        if self.retry_jitter_min_ms is not None and self.retry_jitter_max_ms is not None:
            low = min(self.retry_jitter_min_ms, self.retry_jitter_max_ms)
            high = max(self.retry_jitter_min_ms, self.retry_jitter_max_ms)
            return random.uniform(low, high) / 1000.0
        return backoff

    def _generate_with_retry(self, content: str, account_id: Optional[str]) -> Any:
        """Call generate_content with retry logic for transient errors."""
        last_exc: Optional[Exception] = None
        backoff = INITIAL_BACKOFF_S
        start = time.monotonic()
        total_attempts = self.max_retries + 1
        budget_s = (
            float(self.wall_clock_budget_s)
            if self.wall_clock_budget_s is not None
            else float(self.timeout) * total_attempts
        )

        for attempt in range(total_attempts):
            elapsed = time.monotonic() - start
            remaining = budget_s - elapsed
            if remaining <= 0:
                break
            try:
                timeout_s = min(self.timeout, remaining)
                with LLMGuard(
                    model=self.model_name,
                    account_id=account_id,
                    timeout_s=timeout_s,
                ):
                    response = self._call_with_timeout(content, timeout_s)
                record_success(self.model_name)
                return response
            except (LLMConcurrencyError, LLMCircuitOpenError):
                raise
            except Exception as exc:
                last_exc = exc
                retryable = self._is_retryable_error(exc)
                if not retryable:
                    raise
                if attempt >= self.max_retries:
                    record_failure(self.model_name)
                    raise
                sleep_s = self._retry_sleep_s(backoff)
                logger.warning(
                    "Gemini API error (attempt %d/%d), retrying in %.2fs: %s",
                    attempt + 1,
                    total_attempts,
                    sleep_s,
                    type(exc).__name__,
                )
                time.sleep(sleep_s)
                if self.retry_jitter_min_ms is None:
                    backoff = min(backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_S)

        # All retries exhausted
        if last_exc:
            record_failure(self.model_name)
            raise last_exc
        raise RuntimeError("Unexpected retry loop exit")

    def chat_json(
        self,
        messages: List[Dict[str, Any]],
        retry_strict: bool = True,
        account_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Call Gemini's chat API expecting a JSON response payload.

        Args:
            messages: List of {"role": "user"|"system", "content": str}
            retry_strict: If True, retry once with stricter prompt on parse failure

        Returns:
            Parsed JSON dict from model response

        Raises:
            ValueError: If response is not valid JSON after retries
            Exception: On API errors
        """
        # Convert messages to Gemini format
        gemini_messages = self._convert_messages(messages)

        try:
            # Generate response with retry logic for transient errors
            response = self._generate_with_retry(
                gemini_messages, account_id or self.account_id
            )

            # Extract text
            content_text = response.text.strip()

            if not content_text:
                return {"_parse_error": True, "_raw": "(empty response)"}

            # Parse JSON
            try:
                return self._parse_json_response(content_text)
            except (json.JSONDecodeError, ValueError) as parse_err:
                logger.warning(
                    "Gemini JSON parse failed (first pass). Raw=%s",
                    content_text[:2000],
                )
                # Retry with stricter prompt if enabled
                if retry_strict:
                    return self._retry_with_strict_prompt(messages, content_text)
                else:
                    return {
                        "_parse_error": True,
                        "_raw": content_text,
                        "_error": str(parse_err)
                    }

        except Exception as api_err:
            # Log and return sanitized error
            sanitized = self._sanitize_error(api_err)
            logger.error("Gemini API error after retries: %s", sanitized)
            err_type = type(api_err).__name__
            retry_after = None
            if isinstance(api_err, (LLMConcurrencyError, LLMCircuitOpenError)):
                retry_after = getattr(api_err, "retry_after_s", None)
            return {
                "_parse_error": True,
                "_raw": sanitized,
                "_error": f"API error: {err_type}",
                "_error_type": err_type,
                "_retry_after_s": retry_after,
            }

    def _convert_messages(
        self, messages: List[Dict[str, Any]]
    ) -> str:
        """
        Convert OpenAI-style messages to Gemini format (simple string for Vertex AI).

        Gemini doesn't have a "system" role, so we prepend system messages
        to the first user message.
        """
        system_parts = []
        user_parts = []

        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")

            if role == "system":
                system_parts.append(content)
            else:
                user_parts.append(content)

        # Combine system + user messages into a single string
        if system_parts:
            full_content = "\n\n".join(system_parts) + "\n\n" + "\n\n".join(user_parts)
        else:
            full_content = "\n\n".join(user_parts)

        return full_content

    def _parse_json_response(self, content: str) -> Dict[str, Any]:
        """
        Parse JSON from response, handling common LLM quirks.

        - Strips markdown code fences (```json ... ```)
        - Handles trailing commas
        - Cleans control characters
        """
        # Remove invalid control characters that break JSON parsing
        content = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f]", "", content)

        # Normalize smart quotes to straight quotes
        content = (
            content.replace("\u201c", '"')
            .replace("\u201d", '"')
            .replace("\u2018", "'")
            .replace("\u2019", "'")
        )

        # Remove markdown code fences
        content = re.sub(r"^```(?:json)?\s*", "", content, flags=re.MULTILINE)
        content = re.sub(r"\s*```\s*$", "", content, flags=re.MULTILINE)
        content = content.strip()

        # Try to find JSON object/array bounds
        if not content.startswith(("{", "[")):
            # Search for first { or [
            match = re.search(r"[{\[]", content)
            if match:
                content = content[match.start():]

        # Trim anything after the final closing brace/bracket
        end_idx = max(content.rfind("}"), content.rfind("]"))
        if end_idx != -1:
            content = content[: end_idx + 1]

        # Extract the first balanced JSON object/array if multiple are present
        content = self._extract_first_json(content)

        # Escape unescaped control chars/newlines inside JSON strings
        content = self._escape_control_chars_in_strings(content)
        # Remove stray lines that are just a quote (common LLM glitch)
        content = re.sub(r'^\s*"\s*$', "", content, flags=re.MULTILINE)

        # Parse
        try:
            return json.loads(content)
        except json.JSONDecodeError as parse_err:
            logger.debug(
                "Gemini JSON decode error: %s at line %s col %s",
                parse_err.msg,
                getattr(parse_err, "lineno", "?"),
                getattr(parse_err, "colno", "?"),
            )
            # Try fixing trailing commas (common LLM error)
            fixed = re.sub(r",\s*([}\]])", r"\1", content)
            try:
                return json.loads(fixed)
            except json.JSONDecodeError as parse_err2:
                logger.debug(
                    "Gemini JSON decode error after comma fix: %s at line %s col %s",
                    parse_err2.msg,
                    getattr(parse_err2, "lineno", "?"),
                    getattr(parse_err2, "colno", "?"),
                )
                # As a last resort, try parsing with strict=False (if available)
                try:
                    return json.loads(fixed, strict=False)  # type: ignore[arg-type]
                except Exception:
                    logger.debug("Gemini JSON parse failed: %s", fixed[:200])
                    raise

    @staticmethod
    def _extract_first_json(content: str) -> str:
        """Return the first balanced JSON object/array from content."""
        start_idx = None
        depth = 0
        in_str = False
        escape = False
        for i, ch in enumerate(content):
            if start_idx is None:
                if ch in "{[":
                    start_idx = i
                    depth = 1
                continue
            if escape:
                escape = False
                continue
            if ch == "\\":
                escape = True
                continue
            if ch == '"':
                in_str = not in_str
                continue
            if in_str:
                continue
            if ch in "{[":
                depth += 1
            elif ch in "}]":
                depth -= 1
                if depth == 0:
                    return content[start_idx : i + 1]
        return content

    @staticmethod
    def _escape_control_chars_in_strings(content: str) -> str:
        """Escape newlines/tabs within JSON string literals."""
        out = []
        in_str = False
        escape = False
        for ch in content:
            if escape:
                out.append(ch)
                escape = False
                continue
            if ch == "\\":
                out.append(ch)
                escape = True
                continue
            if ch == '"':
                in_str = not in_str
                out.append(ch)
                continue
            if in_str:
                if ch == "\n":
                    out.append("\\n")
                    continue
                if ch == "\r":
                    out.append("\\r")
                    continue
                if ch == "\t":
                    out.append("\\t")
                    continue
            out.append(ch)
        return "".join(out)

    def _retry_with_strict_prompt(
        self,
        original_messages: List[Dict[str, Any]],
        failed_response: str
    ) -> Dict[str, Any]:
        """Retry with a stricter JSON prompt."""
        strict_suffix = (
            "\n\nIMPORTANT: Respond with ONLY valid JSON. "
            "No markdown, no explanation, no code fences. "
            "Just pure JSON starting with { or [. "
            "Output MUST be minified on a single line (no newlines). "
            "Do NOT include unescaped quotes or newlines inside string values."
        )

        # Append strict instruction to last message
        retry_messages = original_messages.copy()
        if retry_messages:
            retry_messages[-1]["content"] += strict_suffix

        try:
            gemini_messages = self._convert_messages(retry_messages)
            # Use retry wrapper for transient errors
            response = self._generate_with_retry(gemini_messages)
            content_text = response.text.strip()
            return self._parse_json_response(content_text)
        except Exception as retry_err:
            sanitized = self._sanitize_error(retry_err)
            logger.warning(
                "Gemini JSON parse failed after retry. Raw=%s",
                failed_response[:2000],
            )
            return {
                "_parse_error": True,
                "_raw": failed_response,
                "_retry_raw": sanitized,
                "_error": "JSON parse failed after retry"
            }
