# api/providers.py

from __future__ import annotations

import json
import uuid
from aiohttp import ClientSession, ClientTimeout
from typing import AsyncGenerator, List, Dict, Any, Union

from api.logger import setup_logger

logger = setup_logger(__name__)

class AmigoChat:
    url = "https://amigochat.io"
    chat_api_endpoint = "https://api.amigochat.io/v1/chat/completions"
    image_api_endpoint = "https://api.amigochat.io/v1/images/generations"
    default_model = 'gpt-4o-mini'

    chat_models = [
        'gpt-4o',
        default_model,
        'o1-preview',
        'o1-mini',
        'meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo',
        'meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo',
        'claude-3-sonnet-20240229',
        'gemini-1.5-pro',
    ]

    image_models = [
        'flux-pro/v1.1',
        'flux-realism',
        'flux-pro',
        'dalle-e-3',
    ]

    models = chat_models + image_models

    model_aliases = {
        "o1": "o1-preview",
        "llama-3.1-405b": "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
        "llama-3.2-90b": "meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo",
        "claude-3.5-sonnet": "claude-3-sonnet-20240229",
        "gemini-pro": "gemini-1.5-pro",
        "flux-pro": "flux-pro/v1.1",
        "dalle-3": "dalle-e-3",
    }

    persona_ids = {
        'gpt-4o': "gpt",
        'gpt-4o-mini': "amigo",
        'o1-preview': "openai-o-one",
        'o1-mini': "openai-o-one-mini",
        'meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo': "llama-three-point-one",
        'meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo': "llama-3-2",
        'claude-3-sonnet-20240229': "claude",
        'gemini-1.5-pro': "gemini-1-5-pro",
        'flux-pro/v1.1': "flux-1-1-pro",
        'flux-realism': "flux-realism",
        'flux-pro': "flux-pro",
        'dalle-e-3': "dalle-three",
    }

    @classmethod
    def get_model(cls, model: str) -> str:
        return cls.model_aliases.get(model, model if model in cls.models else cls.default_model)

    @classmethod
    def get_persona_id(cls, model: str) -> str:
        return cls.persona_ids.get(model, "amigo")

    @classmethod
    async def generate_response(
        cls,
        model: str,
        messages: List[Dict[str, Any]],
        stream: bool = False,
        proxy: str = None,
    ) -> Union[AsyncGenerator[str, None], str]:
        model = cls.get_model(model)
        device_uuid = str(uuid.uuid4())

        headers = {
            "accept": "*/*",
            "accept-language": "en-US,en;q=0.9",
            "authorization": "Bearer",
            "cache-control": "no-cache",
            "content-type": "application/json",
            "origin": cls.url,
            "pragma": "no-cache",
            "priority": "u=1, i",
            "referer": f"{cls.url}/",
            "sec-ch-ua": '"Chromium";v="129", "Not=A?Brand";v="8"',
            "sec-ch-ua-mobile": "?0",
            "sec-ch-ua-platform": '"Linux"',
            "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36",
            "x-device-language": "en-US",
            "x-device-platform": "web",
            "x-device-uuid": device_uuid,
            "x-device-version": "1.0.32"
        }

        async with ClientSession(headers=headers) as session:
            if model in cls.chat_models:
                data = {
                    "messages": messages,
                    "model": model,
                    "personaId": cls.get_persona_id(model),
                    "frequency_penalty": 0,
                    "max_tokens": 4000,
                    "presence_penalty": 0,
                    "stream": stream,
                    "temperature": 0.5,
                    "top_p": 0.95
                }

                timeout = ClientTimeout(total=300)
                try:
                    async with session.post(cls.chat_api_endpoint, json=data, proxy=proxy, timeout=timeout) as response:
                        if response.status not in (200, 201):
                            error_text = await response.text()
                            raise Exception(f"Error {response.status}: {error_text}")

                        if stream:
                            async def stream_content():
                                async for line in response.content:
                                    line = line.decode('utf-8').strip()
                                    if line.startswith('data: '):
                                        if line == 'data: [DONE]':
                                            break
                                        try:
                                            chunk = json.loads(line[6:])
                                            if 'choices' in chunk and len(chunk['choices']) > 0:
                                                choice = chunk['choices'][0]
                                                if 'delta' in choice:
                                                    content = choice['delta'].get('content')
                                                elif 'text' in choice:
                                                    content = choice['text']
                                                else:
                                                    content = None
                                                if content:
                                                    yield content
                                        except json.JSONDecodeError:
                                            pass
                            return stream_content()
                        else:
                            response_data = await response.json()
                            if 'choices' in response_data and len(response_data['choices']) > 0:
                                content = response_data['choices'][0]['message']['content']
                                return content
                            else:
                                return ""
                except Exception as e:
                    logger.error(f"Error during request: {e}")
                    raise
            else:
                # Handle image models or other cases if necessary
                pass