from __future__ import annotations

import re
import random
import string
import uuid
import json
import logging
from aiohttp import ClientSession, ClientError
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Dict, Any, Optional
from datetime import datetime
from fastapi.responses import StreamingResponse

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Custom exception for model not working
class ModelNotWorkingException(Exception):
    def __init__(self, model: str):
        self.model = model
        self.message = f"The model '{model}' is currently not working. Please wait for NiansuhAI to fix this. Thank you for your patience."
        super().__init__(self.message)

# Mock implementations for ImageResponse and to_data_uri
class ImageResponse:
    def __init__(self, url: str, alt: str):
        self.url = url
        self.alt = alt

def to_data_uri(image: Any) -> str:
    return "data:image/png;base64,..."  # Replace with actual base64 data

class AsyncGeneratorProvider:
    pass

class ProviderModelMixin:
    pass

class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
    url = "https://www.blackbox.ai"
    api_endpoint = "https://www.blackbox.ai/api/chat"
    working = True
    supports_stream = True
    supports_system_message = True
    supports_message_history = True
    
    default_model = 'blackbox'
    models = [
        'blackbox',
        'gemini-1.5-flash',
        "llama-3.1-8b",
        'llama-3.1-70b',
        'llama-3.1-405b',
        'ImageGenerationLV45LJp',
        'gpt-4o',
        'gemini-pro',
        'claude-sonnet-3.5',
    ]

    agentMode = {
        'ImageGenerationLV45LJp': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
    }

    trendingAgentMode = {
        "blackbox": {},
        "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
        "llama-3.1-8b": {'mode': True, 'id': "llama-3.1-8b"},
        'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"},
        'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405b"},
    }
    
    userSelectedModel = {
        "gpt-4o": "gpt-4o",
        "gemini-pro": "gemini-pro",
        'claude-sonnet-3.5': "claude-sonnet-3.5",
    }
    
    model_aliases = {
        "gemini-flash": "gemini-1.5-flash",
        "flux": "ImageGenerationLV45LJp",
    }

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

    @classmethod
    async def create_async_generator(
        cls,
        model: str,
        messages: List[Dict[str, str]],
        proxy: Optional[str] = None,
        image: Optional[Any] = None,
        image_name: Optional[str] = None,
        **kwargs
    ) -> Any:
        model = cls.get_model(model)

        if not cls.working or model not in cls.models:
            raise ModelNotWorkingException(model)

        headers = {
            "accept": "*/*",
            "accept-language": "en-US,en;q=0.9",
            "cache-control": "no-cache",
            "content-type": "application/json",
            "origin": cls.url,
            "pragma": "no-cache",
            "referer": f"{cls.url}/",
            "sec-ch-ua": '"Not;A=Brand";v="24", "Chromium";v="128"',
            "sec-ch-ua-mobile": "?0",
            "sec-ch-ua-platform": '"Linux"',
            "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/128.0.0.0 Safari/537.36"
        }

        if model in cls.userSelectedModel:
            prefix = f"@{cls.userSelectedModel[model]}"
            if not messages[0]['content'].startswith(prefix):
                messages[0]['content'] = f"{prefix} {messages[0]['content']}"
        
        async with ClientSession(headers=headers) as session:
            if image is not None:
                messages[-1]["data"] = {
                    "fileText": image_name,
                    "imageBase64": to_data_uri(image)
                }
            
            random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7))

            data = {
                "messages": messages,
                "id": random_id,
                "previewToken": None,
                "userId": None,
                "codeModelMode": True,
                "agentMode": {},
                "trendingAgentMode": {},
                "userSelectedModel": None,
                "userSystemPrompt": None,
                "isMicMode": False,
                "maxTokens": 8192,
                "playgroundTopP": 0.9,
                "playgroundTemperature": 0.5,
                "isChromeExt": False,
                "githubToken": None,
                "clickedAnswer2": False,
                "clickedAnswer3": False,
                "clickedForceWebSearch": False,
                "visitFromDelta": False,
                "mobileClient": False,
                "webSearchMode": False,
            }

            if model in cls.agentMode:
                data["agentMode"] = cls.agentMode[model]
            elif model in cls.trendingAgentMode:
                data["trendingAgentMode"] = cls.trendingAgentMode[model]
            elif model in cls.userSelectedModel:
                data["userSelectedModel"] = cls.userSelectedModel[model]
            
            try:
                async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
                    response.raise_for_status()
                    response_text = await response.text()

                    # Check if the response is empty or just whitespace
                    if not response_text.strip():
                        raise ModelNotWorkingException(model)

                    if model == 'ImageGenerationLV45LJp':
                        url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text)
                        if url_match:
                            image_url = url_match.group(0)
                            yield ImageResponse(image_url, alt=messages[-1]['content'])
                        else:
                            raise Exception("Image URL not found in the response")
                    
                    # Handle normal responses
                    async for chunk in response.content.iter_any():
                        if chunk:
                            decoded_chunk = chunk.decode(errors='ignore')
                            decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk)
                            if decoded_chunk.strip():
                                yield decoded_chunk
                        
            except ClientError as e:
                logger.error(f"HTTP request failed: {e}")
                raise HTTPException(status_code=503, detail="Service is unavailable. Please try again later.")

# FastAPI app setup
app = FastAPI()

class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    model: str
    messages: List[Message]
    stream: Optional[bool] = False  # Add this for streaming

def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
    return {
        "id": f"chatcmpl-{uuid.uuid4()}",
        "object": "chat.completion.chunk",
        "created": int(datetime.now().timestamp()),
        "model": model,
        "choices": [
            {
                "index": 0,
                "delta": {"content": content, "role": "assistant"},
                "finish_reason": finish_reason,
            }
        ],
        "usage": None,
    }

@app.post("/niansuhai/v1/chat/completions")
async def chat_completions(request: ChatRequest):
    # Validate the model
    valid_models = Blackbox.models + list(Blackbox.userSelectedModel.keys()) + list(Blackbox.model_aliases.keys())
    if request.model not in valid_models:
        raise HTTPException(status_code=400, detail=f"Invalid model name: {request.model}. Valid models are: {valid_models}")

    messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]

    try:
        async_generator = Blackbox.create_async_generator(
            model=request.model,
            messages=messages,
            image=None,
            image_name=None
        )
    except ModelNotWorkingException as e:
        raise HTTPException(status_code=503, detail=str(e))

    if request.stream:
        async def generate():
            async for chunk in async_generator:
                if isinstance(chunk, ImageResponse):
                    image_markdown = f"![image]({chunk.url})"
                    yield f"data: {json.dumps(create_response(image_markdown, request.model))}\n\n"
                else:
                    yield f"data: {json.dumps(create_response(chunk, request.model))}\n\n"
            yield "data: [DONE]\n\n"

        return StreamingResponse(generate(), media_type="text/event-stream")
    else:
        response_content = ""
        async for chunk in async_generator:
            if isinstance(chunk, ImageResponse):
                response_content += f"![image]({chunk.url})\n"
            else:
                response_content += chunk

        return {
            "id": f"chatcmpl-{uuid.uuid4()}",
            "object": "chat.completion",
            "created": int(datetime.now().timestamp()),
            "model": request.model,
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "content": response_content
                    },
                    "finish_reason": "stop",
                    "index": 0
                }
            ],
            "usage": None,
        }

@app.get("/niansuhai/v1/models")
async def get_models():
    return {"models": Blackbox.models}