import os
import re
import json
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from langchain_community.llms import Ollama
import logging
from typing import List, Dict, Optional
from functools import lru_cache
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

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

app = FastAPI()
MODEL_NAME = 'gemma3:1b'

@lru_cache()
def get_llm():
    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
    return Ollama(model=MODEL_NAME, callback_manager=callback_manager)

# OpenAI-compatible models
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    model: str = MODEL_NAME
    messages: List[ChatMessage]
    stream: bool = False
    temperature: Optional[float] = 0.7
    max_tokens: Optional[int] = None

@app.get("/")
def read_root():
    return {"Hello": f"Welcome to {MODEL_NAME} FastAPI"}

@app.get("/models")
async def get_models():
    return {
        "object": "list",
        "data": [{
            "id": MODEL_NAME,
            "object": "model",
            "created": 1677610602,
            "owned_by": "open-webui"
        }]
    }

@app.post("/chat/completions")
async def chat_completion(request: ChatCompletionRequest):
    try:
        # Extract the last user message
        user_messages = [m for m in request.messages if m.role == 'user']
        if not user_messages:
            raise HTTPException(status_code=400, detail="No user message found")
        
        prompt = user_messages[-1].content
        logger.info(f"Received request: {prompt}")

        llm = get_llm()

        if request.stream:
            async def generate():
                full_response = ""
                async for chunk in llm.astream(prompt):
                    full_response += chunk
                    yield json.dumps({
                        "choices": [{
                            "index": 0,
                            "delta": {"content": chunk},
                            "finish_reason": None
                        }]
                    }) + "\n"
                
                # Final completion message
                yield json.dumps({
                    "choices": [{
                        "index": 0,
                        "delta": {},
                        "finish_reason": "stop"
                    }]
                }) + "\n"
                logger.info(f"Full streamed response: {full_response}")

            return StreamingResponse(
                generate(),
                media_type="text/event-stream",
                headers={"Cache-Control": "no-cache"}
            )
        else:
            response = llm.invoke(prompt)
            return {
                "id": "chatcmpl-123",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": MODEL_NAME,
                "choices": [{
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": response
                    },
                    "finish_reason": "stop"
                }],
                "usage": {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0
                }
            }

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

@app.on_event("startup")
async def startup_event():
    logger.info(f"Starting up with model: {MODEL_NAME}")
    get_llm()  # Warm up the model

@app.on_event("shutdown")
async def shutdown_event():
    logger.info("Shutting down")

ทำให้ใช้งานร่วมกับ open-webui ได้