from __future__ import annotations

import asyncio
import aiohttp
import random
import string
import json
from urllib.parse import urlencode

from aiohttp import ClientSession

from ..typing import AsyncResult, Messages
from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
from .helper import format_prompt


class RubiksAI(AsyncGeneratorProvider, ProviderModelMixin):
    label = "Rubiks AI"
    url = "https://rubiks.ai"
    api_endpoint = "https://rubiks.ai/search/api.php"
    working = True
    supports_stream = True
    supports_system_message = True
    supports_message_history = True

    default_model = 'llama-3.1-70b-versatile'
    models = [default_model, 'gpt-4o-mini']

    model_aliases = {
        "llama-3.1-70b": "llama-3.1-70b-versatile",
    }

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

    @staticmethod
    def generate_mid() -> str:
        """
        Generates a 'mid' string following the pattern:
        6 characters - 4 characters - 4 characters - 4 characters - 12 characters
        Example: 0r7v7b-quw4-kdy3-rvdu-ekief6xbuuq4
        """
        parts = [
            ''.join(random.choices(string.ascii_lowercase + string.digits, k=6)),
            ''.join(random.choices(string.ascii_lowercase + string.digits, k=4)),
            ''.join(random.choices(string.ascii_lowercase + string.digits, k=4)),
            ''.join(random.choices(string.ascii_lowercase + string.digits, k=4)),
            ''.join(random.choices(string.ascii_lowercase + string.digits, k=12))
        ]
        return '-'.join(parts)

    @staticmethod
    def create_referer(q: str, mid: str, model: str = '') -> str:
        """
        Creates a Referer URL with dynamic q and mid values, using urlencode for safe parameter encoding.
        """
        params = {'q': q, 'model': model, 'mid': mid}
        encoded_params = urlencode(params)
        return f'https://rubiks.ai/search/?{encoded_params}'

    @classmethod
    async def create_async_generator(
        cls,
        model: str,
        messages: Messages,
        proxy: str = None,
        websearch: bool = False,
        **kwargs
    ) -> AsyncResult:
        """
        Creates an asynchronous generator that sends requests to the Rubiks AI API and yields the response.

        Parameters:
        - model (str): The model to use in the request.
        - messages (Messages): The messages to send as a prompt.
        - proxy (str, optional): Proxy URL, if needed.
        - websearch (bool, optional): Indicates whether to include search sources in the response. Defaults to False.
        """
        model = cls.get_model(model)
        prompt = format_prompt(messages)
        q_value = prompt
        mid_value = cls.generate_mid()
        referer = cls.create_referer(q=q_value, mid=mid_value, model=model)

        url = cls.api_endpoint
        params = {
            'q': q_value,
            'model': model,
            'id': '',
            'mid': mid_value
        }

        headers = {
            'Accept': 'text/event-stream',
            'Accept-Language': 'en-US,en;q=0.9',
            'Cache-Control': 'no-cache',
            'Connection': 'keep-alive',
            'Pragma': 'no-cache',
            'Referer': referer,
            'Sec-Fetch-Dest': 'empty',
            'Sec-Fetch-Mode': 'cors',
            'Sec-Fetch-Site': 'same-origin',
            'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36',
            'sec-ch-ua': '"Chromium";v="129", "Not=A?Brand";v="8"',
            'sec-ch-ua-mobile': '?0',
            'sec-ch-ua-platform': '"Linux"'
        }

        try:
            timeout = aiohttp.ClientTimeout(total=None)
            async with ClientSession(timeout=timeout) as session:
                async with session.get(url, headers=headers, params=params, proxy=proxy) as response:
                    if response.status != 200:
                        yield f"Request ended with status code {response.status}"
                        return

                    assistant_text = ''
                    sources = []

                    async for line in response.content:
                        decoded_line = line.decode('utf-8').strip()
                        if not decoded_line.startswith('data: '):
                            continue
                        data = decoded_line[6:]
                        if data in ('[DONE]', '{"done": ""}'):
                            break
                        try:
                            json_data = json.loads(data)
                        except json.JSONDecodeError:
                            continue

                        if 'url' in json_data and 'title' in json_data:
                            if websearch:
                                sources.append({'title': json_data['title'], 'url': json_data['url']})

                        elif 'choices' in json_data:
                            for choice in json_data['choices']:
                                delta = choice.get('delta', {})
                                content = delta.get('content', '')
                                role = delta.get('role', '')
                                if role == 'assistant':
                                    continue
                                assistant_text += content

                    if websearch and sources:
                        sources_text = '\n'.join([f"{i+1}. [{s['title']}]: {s['url']}" for i, s in enumerate(sources)])
                        assistant_text += f"\n\n**Source:**\n{sources_text}"

                    yield assistant_text

        except asyncio.CancelledError:
            yield "The request was cancelled."
        except aiohttp.ClientError as e:
            yield f"An error occurred during the request: {e}"
        except Exception as e:
            yield f"An unexpected error occurred: {e}"