import os
import json
import time
import logging
import asyncio
import g4f

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


# ================= Logging =================

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s"
)

logger=logging.getLogger(__name__)


# ================= App =================

app=FastAPI(
    title="G4F Smart Gateway",
    version="2.0"
)

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

API_KEY=os.getenv(
    "API_KEY",
    "your_fallback_secret"
)


# ================= Models =================

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


class ChatRequest(BaseModel):
    model:str="gpt-4o-mini"
    messages:List[ChatMessage]
    stream:bool=False
    provider:Optional[str]=None
    temperature:Optional[float]=0.7
    max_tokens:Optional[int]=2048


# ================= Auth =================

def verify_key(auth:str):

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

    parts=auth.split()

    if len(parts)!=2:
        raise HTTPException(
            status_code=401,
            detail="Malformed Authorization"
        )

    if parts[0]!="Bearer":
        raise HTTPException(
            status_code=401,
            detail="Invalid Authorization"
        )

    if parts[1]!=API_KEY:
        raise HTTPException(
            status_code=401,
            detail="Invalid API Key"
        )


# ================= Provider Selection =================

def choose_provider(model:str):

    model=model.lower()

    try:

        if "qwen" in model:

            return getattr(
                g4f.Provider,
                "DeepInfra",
                None
            )

        elif "perplexity" in model:

            return getattr(
                g4f.Provider,
                "PerplexityLabs",
                None
            )

        elif "llama" in model:

            return getattr(
                g4f.Provider,
                "DeepInfra",
                None
            )

        elif "claude" in model:

            return getattr(
                g4f.Provider,
                "OpenaiChat",
                None
            )

        elif "gemini" in model:

            return getattr(
                g4f.Provider,
                "OpenaiChat",
                None
            )

        elif "gpt" in model:

            return getattr(
                g4f.Provider,
                "OpenaiChat",
                None
            )

    except:
        pass

    return None


# ================= Health =================

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

    return {
        "status":"online"
    }


# ================= Models =================

@app.get("/v1/models")
async def models(
    authorization:str=Header(None)
):

    verify_key(
        authorization
    )

    try:

        found=[]

        if hasattr(
            g4f.models,
            "ModelUtils"
        ):

            found=list(
                g4f.models.ModelUtils.convert.keys()
            )

        found=sorted(
            list(set(found))
        )

        if not found:

            found=[

                "gpt-4o-mini",
                "gpt-4",
                "gpt-3.5-turbo",
                "qwen-2.5-72b",
                "llama-3-70b",
                "perplexity"

            ]

        return {

            "object":"list",

            "data":[

                {
                    "id":m,
                    "object":"model",
                    "owned_by":"g4f"
                }

                for m in found

            ]
        }

    except Exception as e:

        logger.exception(e)

        return {

            "object":"list",

            "data":[

                {
                    "id":"gpt-4o-mini",
                    "object":"model"
                }

            ]
        }


# ================= Chat =================

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

    verify_key(
        authorization
    )

    try:

        messages=[

            {
                "role":m.role,
                "content":m.content
            }

            for m in body.messages
        ]


        provider=None

        if body.provider:

            provider=getattr(
                g4f.Provider,
                body.provider,
                None
            )

        else:

            provider=choose_provider(
                body.model
            )


        logger.info(
            f"Model={body.model} Provider={provider}"
        )


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

        if body.stream:

            def generate():

                try:

                    response=g4f.ChatCompletion.create(

                        model=body.model,
                        messages=messages,
                        provider=provider,
                        stream=True

                    )

                    for chunk in response:

                        payload={

                            "id":"chatcmpl",

                            "object":
                            "chat.completion.chunk",

                            "created":
                            int(time.time()),

                            "model":
                            body.model,

                            "choices":[
                                {
                                    "delta":{
                                        "content":
                                        str(chunk)
                                    },
                                    "index":0
                                }
                            ]
                        }

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

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

                except Exception as e:

                    logger.exception(e)

                    yield (
                        f"data:{json.dumps({'error':str(e)})}\n\n"
                    )


            return StreamingResponse(
                generate(),
                media_type="text/event-stream"
            )


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

        response=await asyncio.to_thread(

            g4f.ChatCompletion.create,

            model=body.model,
            messages=messages,
            provider=provider

        )


        return {

            "id":"chatcmpl",

            "object":"chat.completion",

            "created":
            int(time.time()),

            "model":
            body.model,

            "choices":[

                {

                    "index":0,

                    "message":{

                        "role":
                        "assistant",

                        "content":
                        str(response)

                    },

                    "finish_reason":
                    "stop"

                }

            ]
        }

    except Exception as e:

        logger.exception(e)

        raise HTTPException(
            status_code=500,
            detail=str(e)
        )


# ================= Test =================

@app.get("/test/{model}")
async def test_model(
    model:str
):

    try:

        provider=choose_provider(
            model
        )

        result=await asyncio.to_thread(

            g4f.ChatCompletion.create,

            model=model,

            provider=provider,

            messages=[

                {
                    "role":"user",
                    "content":"hello"
                }

            ]
        )

        return {

            "model":model,
            "provider":str(provider),
            "working":True

        }

    except Exception as e:

        return {

            "model":model,
            "working":False,
            "error":str(e)

        }