"""
AI Service - Server-side AI provider integration.

This module handles AI API calls server-side, keeping API keys secure
and providing a unified interface for the frontend.

Environment Variables:
    GEMINI_API_KEY: Google Gemini API key (required for AI features)
    HEXARCH_AI_MODEL: Model to use (default: gemini-2.0-flash)
    HEXARCH_AI_ENABLED: Enable/disable AI features (default: true)
"""

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass
from typing import Any

import httpx

logger = logging.getLogger(__name__)

# Configuration
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
AI_MODEL = os.environ.get("HEXARCH_AI_MODEL", "gemini-2.0-flash")
AI_ENABLED = os.environ.get("HEXARCH_AI_ENABLED", "true").lower() == "true"
GEMINI_API_BASE = "https://generativelanguage.googleapis.com/v1beta"


@dataclass
class AISource:
    """A source reference from AI grounding."""

    title: str
    uri: str


@dataclass
class PolicyConfig:
    """AI-generated policy configuration."""

    name: str
    type: str
    scope: str
    phase: str
    failure_mode: str
    short_circuit: bool
    config: dict[str, Any]


@dataclass
class AIResponse:
    """AI assistant response."""

    text: str
    sources: list[AISource]


def is_ai_available() -> bool:
    """Check if AI service is available."""
    return AI_ENABLED and bool(GEMINI_API_KEY)


def get_ai_status() -> dict[str, Any]:
    """Get AI service status for health checks."""
    return {
        "enabled": AI_ENABLED,
        "configured": bool(GEMINI_API_KEY),
        "model": AI_MODEL if GEMINI_API_KEY else None,
    }


async def generate_policy(prompt: str) -> PolicyConfig:
    """
    Generate a policy configuration using AI.

    Args:
        prompt: Natural language description of the policy

    Returns:
        PolicyConfig with generated settings

    Raises:
        ValueError: If AI is not available or prompt is invalid
        RuntimeError: If AI call fails
    """
    if not is_ai_available():
        raise ValueError("AI service is not available. Check GEMINI_API_KEY configuration.")

    if not prompt or len(prompt.strip()) == 0:
        raise ValueError("Policy prompt cannot be empty")

    if len(prompt) > 5000:
        raise ValueError("Policy prompt exceeds maximum length (5000 characters)")

    system_prompt = """Generate a JSON configuration for a Java-based API Gateway policy.
The system supports:
- Lifecycle states (Draft, Published, Deprecated, Retired)
- Scopes (Global, API, Version, Plan, Route)
- Phases (Pre-Request, Post-Request, Error Handling)
- Failure Modes (Fail Open, Fail Closed)

Return structured JSON with: 'name', 'type', 'scope', 'phase', 'failureMode', 'shortCircuit' (boolean), and a 'config' object specific to the policy type."""

    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.post(
                f"{GEMINI_API_BASE}/models/{AI_MODEL}:generateContent",
                params={"key": GEMINI_API_KEY},
                json={
                    "contents": [
                        {"role": "user", "parts": [{"text": f"{system_prompt}\n\nUser request: {prompt}"}]}
                    ],
                    "generationConfig": {
                        "responseMimeType": "application/json",
                        "temperature": 0.3,
                    },
                },
            )
            response.raise_for_status()
            data = response.json()

        # Extract text from response
        text = data.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0].get("text", "")
        if not text:
            logger.warning("Empty response from Gemini API for policy generation")
            return _fallback_policy()

        # Parse JSON response
        try:
            parsed = json.loads(text)
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse policy JSON: {e}")
            return _fallback_policy()

        return PolicyConfig(
            name=parsed.get("name", "Generated Policy"),
            type=parsed.get("type", "SECURITY"),
            scope=parsed.get("scope", "GLOBAL"),
            phase=parsed.get("phase", "PRE_REQUEST"),
            failure_mode=parsed.get("failureMode", "FAIL_CLOSED"),
            short_circuit=parsed.get("shortCircuit", False),
            config=parsed.get("config", {}),
        )

    except httpx.HTTPStatusError as e:
        logger.error(f"Gemini API error: {e.response.status_code} - {e.response.text}")
        raise RuntimeError(f"AI service error: {e.response.status_code}")
    except Exception as e:
        logger.error(f"Policy generation failed: {e}")
        return _fallback_policy()


async def ask_zenith(query: str) -> AIResponse:
    """
    Query the AI assistant (Zenith AI) with search grounding.

    Args:
        query: Technical question

    Returns:
        AIResponse with text and optional sources

    Raises:
        ValueError: If AI is not available or query is invalid
        RuntimeError: If AI call fails
    """
    if not is_ai_available():
        raise ValueError("AI service is not available. Check GEMINI_API_KEY configuration.")

    if not query or len(query.strip()) == 0:
        raise ValueError("Query cannot be empty")

    if len(query) > 2000:
        raise ValueError("Query exceeds maximum length (2000 characters)")

    system_instruction = """You are Zenith AI, an expert in high-throughput Java API Gateways,
Kubernetes Operators, and protocol mediation (REST/SOAP/MQTT/gRPC).
Provide accurate, technical answers with practical examples."""

    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.post(
                f"{GEMINI_API_BASE}/models/{AI_MODEL}:generateContent",
                params={"key": GEMINI_API_KEY},
                json={
                    "contents": [{"role": "user", "parts": [{"text": query}]}],
                    "systemInstruction": {"parts": [{"text": system_instruction}]},
                    "generationConfig": {
                        "temperature": 0.7,
                    },
                    "tools": [{"googleSearch": {}}],
                },
            )
            response.raise_for_status()
            data = response.json()

        # Extract text from response
        candidate = data.get("candidates", [{}])[0]
        text = candidate.get("content", {}).get("parts", [{}])[0].get("text", "")

        if not text or len(text.strip()) < 10:
            logger.warning("Minimal response from Gemini API")
            return _fallback_response()

        # Extract grounding sources if available
        sources: list[AISource] = []
        grounding_metadata = candidate.get("groundingMetadata", {})
        grounding_chunks = grounding_metadata.get("groundingChunks", [])

        for chunk in grounding_chunks:
            web = chunk.get("web", {})
            uri = web.get("uri")
            if uri:
                sources.append(
                    AISource(
                        title=web.get("title", "Source"),
                        uri=uri,
                    )
                )

        return AIResponse(text=text, sources=sources)

    except httpx.HTTPStatusError as e:
        logger.error(f"Gemini API error: {e.response.status_code} - {e.response.text}")
        raise RuntimeError(f"AI service error: {e.response.status_code}")
    except Exception as e:
        logger.error(f"AI query failed: {e}")
        return _fallback_response()


async def suggest_transformation(source_format: str, target_format: str) -> str:
    """
    Suggest protocol transformation approach.

    Args:
        source_format: Source protocol (e.g., REST, SOAP)
        target_format: Target protocol

    Returns:
        Transformation advice as text

    Raises:
        ValueError: If AI is not available or formats are invalid
        RuntimeError: If AI call fails
    """
    if not is_ai_available():
        raise ValueError("AI service is not available. Check GEMINI_API_KEY configuration.")

    if not source_format or not target_format:
        raise ValueError("Both source and target formats are required")

    prompt = f"""Design a high-performance mediation layer to transform {source_format} to {target_format}.
Explain memory management strategies (e.g. Netty ByteBuf), schema mapping challenges,
and provide a Java-based logic snippet or JSONATA mapping."""

    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.post(
                f"{GEMINI_API_BASE}/models/{AI_MODEL}:generateContent",
                params={"key": GEMINI_API_KEY},
                json={
                    "contents": [{"role": "user", "parts": [{"text": prompt}]}],
                    "generationConfig": {
                        "temperature": 0.5,
                    },
                },
            )
            response.raise_for_status()
            data = response.json()

        text = data.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0].get("text", "")

        if not text or len(text.strip()) == 0:
            return f"Unable to generate transformation advice for {source_format} -> {target_format}"

        return text

    except httpx.HTTPStatusError as e:
        logger.error(f"Gemini API error: {e.response.status_code} - {e.response.text}")
        raise RuntimeError(f"AI service error: {e.response.status_code}")
    except Exception as e:
        logger.error(f"Transformation suggestion failed: {e}")
        return f"Error generating transformation advice. Please try again."


def _fallback_policy() -> PolicyConfig:
    """Return safe fallback policy when AI fails."""
    return PolicyConfig(
        name="Service Temporarily Unavailable",
        type="SECURITY",
        scope="GLOBAL",
        phase="PRE_REQUEST",
        failure_mode="FAIL_CLOSED",
        short_circuit=False,
        config={"note": "AI service is temporarily unavailable. Policy generated with safe defaults."},
    )


def _fallback_response() -> AIResponse:
    """Return fallback response when AI fails."""
    return AIResponse(
        text="Service temporarily unavailable. Please try again in a moment.",
        sources=[],
    )
