import os
import json
import time
import httpx
import logging
from typing import Dict, Any, Optional
from core.metrics import AI_REQUESTS_TOTAL, AI_LATENCY_SECONDS

logger = logging.getLogger(__name__)

class LLMClient:
    def __init__(
        self,
        provider: Optional[str] = None,
        api_key: Optional[str] = None,
        model: Optional[str] = None,
        base_url: Optional[str] = None,
    ):
        self.provider = provider or os.getenv("AI_PROVIDER", "openai").lower()
        self.api_key = api_key or os.getenv("AI_API_KEY", "")
        self.model = model or os.getenv("AI_MODEL", "orcarouter/auto" if self.provider == "openai" else "gemini-1.5-flash")
        self.base_url = base_url or os.getenv("AI_BASE_URL", "https://api.orcarouter.ai/v1")

    async def generate_json(self, prompt: str, system_prompt: Optional[str] = None, action_name: str = "general") -> Dict[str, Any]:
        """Send prompt to LLM and parse JSON response."""
        start_time = time.time()
        status = "error"
        try:
            if self.provider == "gemini" and "orcarouter" not in (self.base_url or ""):
                result = await self._call_gemini(prompt, system_prompt)
            else:
                result = await self._call_openai(prompt, system_prompt)
            status = "success"
            return result
        except Exception as e:
            logger.error(f"LLM generation failed ({self.provider}/{self.model}): {e}")
            raise
        finally:
            duration = time.time() - start_time
            AI_LATENCY_SECONDS.labels(action=action_name).observe(duration)
            AI_REQUESTS_TOTAL.labels(action=action_name, status=status).inc()

    async def _call_gemini(self, prompt: str, system_prompt: Optional[str] = None) -> Dict[str, Any]:
        url = f"https://generativelanguage.googleapis.com/v1beta/models/{self.model}:generateContent?key={self.api_key}"
        payload: Dict[str, Any] = {
            "contents": [
                {
                    "parts": [{"text": prompt}]
                }
            ],
            "generationConfig": {
                "responseMimeType": "application/json",
                "temperature": 0.2,
            }
        }
        if system_prompt:
            payload["systemInstruction"] = {
                "parts": [{"text": system_prompt}]
            }

        async with httpx.AsyncClient(timeout=60.0) as client:
            resp = await client.post(url, json=payload)
            resp.raise_for_status()
            data = resp.json()
            raw_text = data["candidates"][0]["content"]["parts"][0]["text"]
            return json.loads(raw_text)

    async def _call_openai(self, prompt: str, system_prompt: Optional[str] = None) -> Dict[str, Any]:
        base = (self.base_url or "https://api.orcarouter.ai/v1").rstrip("/")
        url = base if base.endswith("/chat/completions") else f"{base}/chat/completions"

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        payload = {
            "model": self.model,
            "messages": messages,
            "response_format": {"type": "json_object"},
            "temperature": 0.2,
        }

        async with httpx.AsyncClient(timeout=60.0) as client:
            resp = await client.post(url, headers=headers, json=payload)
            resp.raise_for_status()
            data = resp.json()
            content = data["choices"][0]["message"]["content"]
            return json.loads(content)