"""
Gemini API Fallback for High-Se Edge Queries (WO-EDGE-DEPTH-ROUTER-01)

When a query's Symbolic Entropy (Se) is too high for local inference,
or when available RAM is insufficient, route to Gemini API via Google AI Studio.

This provides a cloud fallback for edge devices that can't handle
complex queries locally.

Environment:
    GEMINI_API_KEY: Required. Get from https://aistudio.google.com/apikey

Usage:
    from helix_cdc.regrow.gemini_fallback import query_gemini

    result = query_gemini("Explain quantum entanglement")
    print(result.response_text)
"""

import os
import time
import hashlib
from dataclasses import dataclass, field
from typing import Optional, Dict, Any, List

# Use httpx for simple HTTP (no heavy dependencies)
try:
    import httpx
    HTTPX_AVAILABLE = True
except ImportError:
    HTTPX_AVAILABLE = False

# Fallback to urllib if httpx not available
import urllib.request
import urllib.error
import json


# =============================================================================
# Constants
# =============================================================================

GEMINI_API_BASE = "https://generativelanguage.googleapis.com/v1beta"
GEMINI_MODEL = "gemini-2.0-flash"  # Fast, capable model for edge fallback
GEMINI_TIMEOUT_SECONDS = 30.0

# Safety settings - permissive for general use
DEFAULT_SAFETY_SETTINGS = [
    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
]

# Generation config
DEFAULT_GENERATION_CONFIG = {
    "temperature": 0.7,
    "topP": 0.95,
    "topK": 40,
    "maxOutputTokens": 2048,
}


# =============================================================================
# GeminiFallbackResult Dataclass
# =============================================================================

@dataclass
class GeminiFallbackResult:
    """
    Result from Gemini API fallback query.

    Contains the response, metadata, and diagnostic information
    for receipts and debugging.
    """
    # Response
    response_text: str
    success: bool

    # Error handling
    error_message: Optional[str] = None
    error_code: Optional[str] = None

    # Request metadata
    prompt: str = ""
    prompt_hash: str = ""
    model: str = GEMINI_MODEL

    # Timing
    request_time_ms: float = 0.0
    total_time_ms: float = 0.0

    # Token counts (if available from API)
    prompt_tokens: Optional[int] = None
    completion_tokens: Optional[int] = None

    # Receipt schema
    schema: str = "gemini_fallback_result_v1"
    work_order: str = "WO-EDGE-DEPTH-ROUTER-01"

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        return {
            "schema": self.schema,
            "work_order": self.work_order,
            "success": self.success,
            "response_text": self.response_text,
            "error_message": self.error_message,
            "error_code": self.error_code,
            "prompt_hash": self.prompt_hash,
            "model": self.model,
            "request_time_ms": self.request_time_ms,
            "total_time_ms": self.total_time_ms,
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
        }


# =============================================================================
# Core Query Function
# =============================================================================

def query_gemini(
    prompt: str,
    api_key: Optional[str] = None,
    model: str = GEMINI_MODEL,
    temperature: float = 0.7,
    max_tokens: int = 2048,
    timeout_seconds: float = GEMINI_TIMEOUT_SECONDS,
    system_prompt: Optional[str] = None,
) -> GeminiFallbackResult:
    """
    Query Gemini API as fallback for high-Se edge queries.

    This function provides a simple, reliable interface to Gemini API
    for when local inference isn't feasible (high Se or low RAM).

    Args:
        prompt: The user prompt to send to Gemini.
        api_key: Gemini API key. If None, reads from GEMINI_API_KEY env var.
        model: Gemini model to use (default: gemini-2.0-flash).
        temperature: Sampling temperature (0.0 to 1.0).
        max_tokens: Maximum output tokens.
        timeout_seconds: Request timeout.
        system_prompt: Optional system prompt to prepend.

    Returns:
        GeminiFallbackResult with response or error details.

    Raises:
        ValueError: If no API key is available.

    Examples:
        >>> result = query_gemini("What is 2+2?")
        >>> if result.success:
        ...     print(result.response_text)
        ... else:
        ...     print(f"Error: {result.error_message}")
    """
    start_time = time.perf_counter()

    # Get API key
    api_key = api_key or os.environ.get("GEMINI_API_KEY")
    if not api_key:
        return GeminiFallbackResult(
            response_text="",
            success=False,
            error_message="No API key: Set GEMINI_API_KEY environment variable",
            error_code="NO_API_KEY",
            prompt=prompt,
            prompt_hash=_hash_prompt(prompt),
            model=model,
            total_time_ms=(time.perf_counter() - start_time) * 1000,
        )

    # Build request
    url = f"{GEMINI_API_BASE}/models/{model}:generateContent?key={api_key}"

    # Build content parts
    parts = []
    if system_prompt:
        parts.append({"text": f"System: {system_prompt}\n\n"})
    parts.append({"text": prompt})

    request_body = {
        "contents": [{"parts": parts}],
        "generationConfig": {
            "temperature": temperature,
            "topP": 0.95,
            "topK": 40,
            "maxOutputTokens": max_tokens,
        },
        "safetySettings": DEFAULT_SAFETY_SETTINGS,
    }

    # Make request
    request_start = time.perf_counter()

    try:
        if HTTPX_AVAILABLE:
            response_data = _query_with_httpx(url, request_body, timeout_seconds)
        else:
            response_data = _query_with_urllib(url, request_body, timeout_seconds)
    except Exception as e:
        return GeminiFallbackResult(
            response_text="",
            success=False,
            error_message=str(e),
            error_code="REQUEST_FAILED",
            prompt=prompt,
            prompt_hash=_hash_prompt(prompt),
            model=model,
            request_time_ms=(time.perf_counter() - request_start) * 1000,
            total_time_ms=(time.perf_counter() - start_time) * 1000,
        )

    request_time_ms = (time.perf_counter() - request_start) * 1000

    # Parse response
    try:
        response_text, prompt_tokens, completion_tokens = _parse_gemini_response(
            response_data
        )
    except Exception as e:
        return GeminiFallbackResult(
            response_text="",
            success=False,
            error_message=f"Failed to parse response: {e}",
            error_code="PARSE_FAILED",
            prompt=prompt,
            prompt_hash=_hash_prompt(prompt),
            model=model,
            request_time_ms=request_time_ms,
            total_time_ms=(time.perf_counter() - start_time) * 1000,
        )

    return GeminiFallbackResult(
        response_text=response_text,
        success=True,
        prompt=prompt,
        prompt_hash=_hash_prompt(prompt),
        model=model,
        request_time_ms=request_time_ms,
        total_time_ms=(time.perf_counter() - start_time) * 1000,
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
    )


# =============================================================================
# HTTP Implementations
# =============================================================================

def _query_with_httpx(
    url: str,
    request_body: Dict[str, Any],
    timeout_seconds: float,
) -> Dict[str, Any]:
    """Make request using httpx (preferred)."""
    with httpx.Client(timeout=timeout_seconds) as client:
        response = client.post(
            url,
            json=request_body,
            headers={"Content-Type": "application/json"},
        )
        response.raise_for_status()
        return response.json()


def _query_with_urllib(
    url: str,
    request_body: Dict[str, Any],
    timeout_seconds: float,
) -> Dict[str, Any]:
    """Make request using urllib (fallback)."""
    data = json.dumps(request_body).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    with urllib.request.urlopen(req, timeout=timeout_seconds) as response:
        return json.loads(response.read().decode("utf-8"))


def _parse_gemini_response(
    response_data: Dict[str, Any],
) -> tuple[str, Optional[int], Optional[int]]:
    """
    Parse Gemini API response.

    Returns:
        Tuple of (response_text, prompt_tokens, completion_tokens).
    """
    # Check for error
    if "error" in response_data:
        error = response_data["error"]
        raise ValueError(f"API error: {error.get('message', 'Unknown error')}")

    # Extract response text
    candidates = response_data.get("candidates", [])
    if not candidates:
        raise ValueError("No candidates in response")

    content = candidates[0].get("content", {})
    parts = content.get("parts", [])
    if not parts:
        raise ValueError("No parts in response content")

    response_text = parts[0].get("text", "")

    # Extract token counts if available
    usage = response_data.get("usageMetadata", {})
    prompt_tokens = usage.get("promptTokenCount")
    completion_tokens = usage.get("candidatesTokenCount")

    return response_text, prompt_tokens, completion_tokens


# =============================================================================
# Utilities
# =============================================================================

def _hash_prompt(prompt: str) -> str:
    """Generate SHA256 hash of prompt for receipts."""
    return hashlib.sha256(prompt.encode("utf-8")).hexdigest()[:16]


def check_gemini_available() -> tuple[bool, str]:
    """
    Check if Gemini API is available and configured.

    Returns:
        Tuple of (available, message).
    """
    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        return False, "GEMINI_API_KEY environment variable not set"

    if len(api_key) < 20:
        return False, "GEMINI_API_KEY appears invalid (too short)"

    return True, f"Gemini API configured (key: {api_key[:8]}...)"


def get_gemini_models() -> List[str]:
    """
    List available Gemini models for edge fallback.

    Returns recommended models in order of preference:
    - gemini-2.0-flash: Fast, good for edge (default)
    - gemini-1.5-flash: Faster, lower quality
    - gemini-1.5-pro: Higher quality, slower
    """
    return [
        "gemini-2.0-flash",
        "gemini-1.5-flash",
        "gemini-1.5-pro",
    ]
