import json
import re
import time
from typing import Optional

from .models import AgentConfig


class AgentService:

    @classmethod
    def execute_prompt(
        cls,
        config: AgentConfig,
        prompt: str,
        history: Optional[list] = None,
        system_prompt_override: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
    ):
        if not config or not config.is_active:
            return False, "Agent not found or inactive", None

        system = system_prompt_override or config.system_prompt or ""
        temp = temperature if temperature is not None else config.temperature
        tokens = max_tokens if max_tokens is not None else config.max_tokens

        messages = [{"role": "system", "content": system}]
        if history:
            for msg in history[-20:]:
                messages.append({"role": msg.get("role", "user"), "content": msg.get("content", "")})
        messages.append({"role": "user", "content": prompt})

        try:
            if config.provider == AgentConfig.Provider.OPENAI:
                response = cls._call_openai(config, messages, temp, tokens)
            elif config.provider == AgentConfig.Provider.DEEPSEEK:
                response = cls._call_deepseek(config, messages, temp, tokens)
            elif config.provider == AgentConfig.Provider.GOOGLE:
                response = cls._call_google(config, messages, temp, tokens)
            elif config.provider == AgentConfig.Provider.OLLAMA:
                response = cls._call_ollama(config, messages, temp, tokens)
            elif config.provider == AgentConfig.Provider.CUSTOM:
                response = cls._call_custom(config, messages, temp, tokens)
            else:
                return False, f"Unsupported provider: {config.provider}", None

            return True, response.get("content", ""), response.get("tokens_used")

        except Exception as e:
            return False, f"Error calling {config.provider}: {str(e)}", None

    @classmethod
    def _call_openai(cls, config, messages, temperature, max_tokens):
        from openai import OpenAI
        client = OpenAI(api_key=config.api_key or None)
        resp = client.chat.completions.create(
            model=config.model_name,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        return {
            "content": resp.choices[0].message.content or "",
            "tokens_used": resp.usage.total_tokens if resp.usage else 0,
        }

    @classmethod
    def _call_deepseek(cls, config, messages, temperature, max_tokens):
        from openai import OpenAI
        client = OpenAI(
            api_key=config.api_key or None,
            base_url=config.base_url or "https://api.deepseek.com",
        )
        resp = client.chat.completions.create(
            model=config.model_name or "deepseek-chat",
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        return {
            "content": resp.choices[0].message.content or "",
            "tokens_used": resp.usage.total_tokens if resp.usage else 0,
        }

    @classmethod
    def _call_google(cls, config, messages, temperature, max_tokens):
        from google import genai
        client = genai.Client(api_key=config.api_key or None)
        system_msg = ""
        chat_messages = []
        for m in messages:
            if m["role"] == "system":
                system_msg += m["content"] + "\n"
            else:
                chat_messages.append({"role": m["role"], "parts": [m["content"]]})
        model = client.models.generate_content(
            model=config.model_name or "gemini-2.0-flash",
            contents=chat_messages,
            config={
                "system_instruction": system_msg.strip() if system_msg else None,
                "temperature": temperature,
                "max_output_tokens": max_tokens,
            },
        )
        return {
            "content": model.text or "",
            "tokens_used": 0,
        }

    @classmethod
    def _call_ollama(cls, config, messages, temperature, max_tokens):
        import requests
        url = (config.base_url or "http://localhost:11434") + "/api/chat"
        payload = {
            "model": config.model_name or "llama3",
            "messages": messages,
            "options": {
                "temperature": temperature,
                "num_predict": max_tokens,
            },
        }
        resp = requests.post(url, json=payload, timeout=120)
        resp.raise_for_status()
        data = resp.json()
        return {
            "content": data.get("message", {}).get("content", ""),
            "tokens_used": 0,
        }

    @classmethod
    def _call_custom(cls, config, messages, temperature, max_tokens):
        from openai import OpenAI
        client = OpenAI(
            api_key=config.api_key or "fake-key",
            base_url=config.base_url or "http://localhost:8000/v1",
        )
        resp = client.chat.completions.create(
            model=config.model_name or "custom-model",
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
        )
        return {
            "content": resp.choices[0].message.content or "",
            "tokens_used": resp.usage.total_tokens if resp.usage else 0,
        }