# api/utils.py

from datetime import datetime
import json
import uuid
from typing import Any, Dict, Optional, AsyncGenerator, List

from fastapi import HTTPException
from api.models import ChatRequest, Message
from api.logger import setup_logger
from api.providers import AmigoChat

logger = setup_logger(__name__)

async def process_streaming_response(request: ChatRequest) -> AsyncGenerator[str, None]:
    logger.info("Processing streaming response with AmigoChat")
    messages = [msg.dict() for msg in request.messages]

    try:
        async for content in AmigoChat.generate_response(
            model=request.model,
            messages=messages,
            stream=True
        ):
            timestamp = int(datetime.now().timestamp())
            chunk = {
                "id": f"chatcmpl-{uuid.uuid4()}",
                "object": "chat.completion.chunk",
                "created": timestamp,
                "model": request.model,
                "choices": [
                    {
                        "index": 0,
                        "delta": {"content": content},
                        "finish_reason": None,
                    }
                ],
            }
            yield f"data: {json.dumps(chunk)}\n\n"

        # Indicate the end of the stream
        end_chunk = {
            "id": f"chatcmpl-{uuid.uuid4()}",
            "object": "chat.completion.chunk",
            "created": int(datetime.now().timestamp()),
            "model": request.model,
            "choices": [
                {
                    "index": 0,
                    "delta": {},
                    "finish_reason": "stop",
                }
            ],
        }
        yield f"data: {json.dumps(end_chunk)}\n\n"
        yield "data: [DONE]\n\n"

    except Exception as e:
        logger.error(f"Error in streaming response: {e}")
        raise HTTPException(status_code=500, detail=str(e))