import httpx

from fastapi import APIRouter, HTTPException, Header, Request, Response
from httpx_sse import EventSource
from sse_starlette.sse import EventSourceResponse
from starlette.background import BackgroundTask
from app.internal.constants import DUCKDUCKGO_CHAT_ENDPOINT, DUCKDUCKGO_STATUS_ENDPOINT, SESSION_KEY
from typing import Annotated, AsyncIterator, List, Literal
from pydantic import BaseModel

DONE = '[DONE]'


class Message(BaseModel):
    role: Literal['assistant', 'user']
    content: str


class Chat(BaseModel):
    model: Literal[
        'gpt-3.5-turbo-0125',
        'claude-3-haiku-20240307'
    ] = 'gpt-3.5-turbo-0125'
    messages: list[Message]
    stream: bool = False

    model_config = {
        'json_schema_extra': {
            'examples': [
                {
                    'model': 'claude-3-haiku-20240307',
                    'messages': [{
                        'role': 'user',
                        'content': 'Hello',
                    }]
                }
            ]
        }
    }


class Choice(BaseModel):
    message: Message


class CompletionsResult(BaseModel):
    choices: List[Choice]


router = APIRouter()


@router.post('/ddg/chat/completions',
             response_model=CompletionsResult,
             responses={
                 200: {
                     'content': {'text/event-stream': {}},
                     'description': 'Return the JSON completions result or an event stream.',
                 }
             })
async def chat(input: Chat, request: Request, response: Response, x_session_id: Annotated[str | None, Header()] = None):
    http_client: httpx.AsyncClient = request.state.http_client
    session_id = x_session_id or (await http_client.get(DUCKDUCKGO_STATUS_ENDPOINT, headers={'x-vqd-accept': '1'})).headers.get(SESSION_KEY)

    req = http_client.build_request('POST', DUCKDUCKGO_CHAT_ENDPOINT,
                                    json=input.model_dump(exclude={'stream'}),
                                    headers={SESSION_KEY: session_id})
    resp = await http_client.send(req, stream=True)

    if resp.status_code != 200:
        raise HTTPException(status_code=400)

    async def agenerator() -> AsyncIterator[str]:
        async for event in EventSource(resp).aiter_sse():
            if event.data == DONE:
                return

            if 'message' in (decoded := event.json()):
                yield decoded['message']

    async def event_generator():
        async for chunk in agenerator():
            yield {
                'data': {
                    'choices': [{
                        'delta': chunk
                    }]
                }
            }

        yield DONE

    response.headers['x-session-id'] = resp.headers.get(SESSION_KEY)

    if input.stream:
        return EventSourceResponse(event_generator(), background=BackgroundTask(resp.aclose), headers=response.headers)

    content = ''
    async for chunk in agenerator():
        content += chunk

    return CompletionsResult(choices=[Choice(message=Message(role='assistant', content=content))])