import asyncio
import json
from concurrent.futures import ThreadPoolExecutor
from typing import List

import g4f

from fastapi import FastAPI, HTTPException, Header
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

# =====================================
# CONFIG
# =====================================

API_KEY = "sk-your-secret-key"

executor = ThreadPoolExecutor(max_workers=5)

# =====================================
# APP
# =====================================

app = FastAPI(
    title="AI Gateway"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# =====================================
# MODELS
# =====================================

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

class ChatRequest(BaseModel):
    model: str = "gpt-4o-mini"
    messages: List[Message]
    stream: bool = False

# =====================================
# AUTH
# =====================================

def verify_api_key(auth: str):

    if not auth:
        raise HTTPException(
            status_code=401,
            detail="Missing Authorization Header"
        )

    if not auth.startswith("Bearer "):
        raise HTTPException(
            status_code=401,
            detail="Invalid Authorization"
        )

    token = auth.replace("Bearer ", "")

    if token != API_KEY:
        raise HTTPException(
            status_code=403,
            detail="Invalid API Key"
        )

# =====================================
# GENERATE
# =====================================

def generate(model, messages):

    response = g4f.ChatCompletion.create(
        model=model,
        messages=messages
    )

    return response

# =====================================
# STREAM
# =====================================

async def stream_generate(model, messages):

    loop = asyncio.get_event_loop()

    response = await loop.run_in_executor(
        executor,
        lambda: g4f.ChatCompletion.create(
            model=model,
            messages=messages,
            stream=True
        )
    )

    for chunk in response:

        if chunk:

            payload = {
                "choices": [
                    {
                        "delta": {
                            "content": chunk
                        }
                    }
                ]
            }

            yield f"data: {json.dumps(payload)}\\n\\n"

    yield "data: [DONE]\\n\\n"

# =====================================
# HOME
# =====================================

@app.get("/")
async def home():

    return {
        "status": "online"
    }

# =====================================
# CHAT
# =====================================

@app.post("/v1/chat/completions")
async def chat(
    req: ChatRequest,
    authorization: str = Header(None)
):

    verify_api_key(authorization)

    messages = [
        m.model_dump()
        for m in req.messages
    ]

    # =========================
    # STREAM
    # =========================

    if req.stream:

        return StreamingResponse(
            stream_generate(
                req.model,
                messages
            ),
            media_type="text/event-stream"
        )

    # =========================
    # NORMAL
    # =========================

    loop = asyncio.get_event_loop()

    response = await loop.run_in_executor(
        executor,
        lambda: generate(
            req.model,
            messages
        )
    )

    return {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "content": response
                }
            }
        ]
    }