app.py

import os
import time
import uuid
from typing import List, Dict, Optional, Union, Generator, Any

from fastapi import FastAPI, HTTPException, Request, status
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel, Field
import uvicorn

from hugchat import hugchat
from hugchat.login import Login
# from hugchat.types.message import MessageNode # For type hinting if needed

# --- Configuration ---
HF_EMAIL = "xawet73334@magpit.com"
HF_PASSWD = "Xawet73334@magpit.com"
COOKIE_PATH_DIR = "./hugchat_cookies/"

if not HF_EMAIL or not HF_PASSWD:
    print("Warning: HUGGINGFACE_EMAIL or HUGGINGFACE_PASSWD environment variables not set.")
    # Allow running without credentials if cookies already exist, for example.
    # The startup logic will handle login/cookie loading.

# --- Global HugChatBot instance and model info ---
chatbot: Optional[hugchat.ChatBot] = None
available_models_list: List[str] = []
available_models_map: Dict[str, int] = {} # Maps model name to index
current_llm_model_on_chatbot: Optional[str] = None
server_start_time = int(time.time()) # For 'created' timestamps

# --- Pydantic Models for OpenAI Compatibility ---

# Model for /v1/models
class ModelCard(BaseModel):
    id: str
    object: str = "model"
    created: int = Field(default_factory=lambda: server_start_time)
    owned_by: str = "huggingface" # Or parse from model ID if possible
    # Add other common fields if desired, often with default/null values
    # permission: Optional[List[Any]] = None
    # root: Optional[str] = None
    # parent: Optional[str] = None

class ModelList(BaseModel):
    object: str = "list"
    data: List[ModelCard]

# Models for /v1/chat/completions (from previous example)
class ChatMessage(BaseModel):
    role: str
    content: str
    # name: Optional[str] = None # For function calling, not directly supported by hugchat

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    stream: Optional[bool] = False
    temperature: Optional[float] = Field(None, ge=0.0, le=2.0) # hugchat might not support all
    top_p: Optional[float] = Field(None, ge=0.0, le=1.0)      # these params directly
    n: Optional[int] = Field(None, ge=1)                      # often n=1 for chat
    max_tokens: Optional[int] = Field(None, ge=1)
    # presence_penalty: Optional[float] = None
    # frequency_penalty: Optional[float] = None
    # logit_bias: Optional[Dict[str, float]] = None
    # user: Optional[str] = None # For tracking, not used by hugchat
    # stop: Optional[Union[str, List[str]]] = None # hugchat handles its own stop

class DeltaMessage(BaseModel):
    role: Optional[str] = None
    content: Optional[str] = None

class ChatCompletionChunkChoice(BaseModel):
    index: int = 0
    delta: DeltaMessage
    finish_reason: Optional[str] = None # "stop", "length", "content_filter", "tool_calls"

class ChatCompletionChunk(BaseModel):
    id: str
    object: str = "chat.completion.chunk"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    # system_fingerprint: Optional[str] = None # OpenAI specific
    choices: List[ChatCompletionChunkChoice]

class ResponseMessage(BaseModel):
    role: str
    content: str
    # tool_calls: Optional[List[Any]] = None # For function/tool calling

class ChatCompletionChoice(BaseModel):
    index: int = 0
    message: ResponseMessage
    finish_reason: str = "stop"
    # logprobs: Optional[Any] = None

class UsageInfo(BaseModel): # Mocked, as hugchat doesn't provide token counts
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0

class ChatCompletionResponse(BaseModel):
    id: str
    object: str = "chat.completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    # system_fingerprint: Optional[str] = None
    choices: List[ChatCompletionChoice]
    usage: Optional[UsageInfo] = Field(default_factory=lambda: UsageInfo())


# --- FastAPI App ---
app = FastAPI(
    title="HugChat OpenAI-Compatible API",
    description="An OpenAI-compatible API wrapper for HuggingChat.",
    version="0.1.1" # Incremented version
)

@app.on_event("startup")
async def startup_event():
    global chatbot, available_models_list, available_models_map, current_llm_model_on_chatbot
    print("Initializing HugChatBot...")
    try:
        if not os.path.exists(COOKIE_PATH_DIR):
            os.makedirs(COOKIE_PATH_DIR)

        if not HF_EMAIL or not HF_PASSWD:
            print("Attempting to load cookies directly as credentials are not fully set.")
            # Try to create a Login object just to access cookie loading methods
            # This part might need adjustment based on how Login handles missing credentials
            temp_sign = Login(HF_EMAIL or "dummy_email", None) # Pass dummy email if HF_EMAIL is None
            cookies = temp_sign.loadCookiesFromDir(cookie_dir_path=COOKIE_PATH_DIR)
            if not cookies:
                raise ValueError("Credentials not set and no saved cookies found. Please set HUGGINGFACE_EMAIL and HUGGINGFACE_PASSWD or ensure cookies are present.")
            print("Loaded cookies from disk.")
        else:
            sign = Login(HF_EMAIL, HF_PASSWD)
            cookies = sign.login(cookie_dir_path=COOKIE_PATH_DIR, save_cookies=True)
        
        chatbot = hugchat.ChatBot(cookies=cookies.get_dict())
        print("HugChatBot initialized successfully.")

        models_raw = chatbot.get_available_llm_models()
        if not models_raw:
            print("Warning: No available LLM models found from HugChat.")
            return

        available_models_list = [str(model_name) for model_name in models_raw]
        available_models_map = {name: i for i, name in enumerate(available_models_list)}
        print(f"Available models: {available_models_list}")

        if available_models_list:
            default_model_index = 0
            chatbot.switch_llm(default_model_index)
            current_llm_model_on_chatbot = available_models_list[default_model_index]
            chatbot.new_conversation(switch_to=True) # Ensure new convo uses this model
            print(f"Default model set to: {current_llm_model_on_chatbot}")
        else:
            print("No models available to set a default.")

    except Exception as e:
        print(f"Error during HugChatBot initialization: {e}")
        chatbot = None

# --- Helper for Unsupported Endpoints ---
def not_supported_response(feature: str):
    return JSONResponse(
        status_code=status.HTTP_501_NOT_IMPLEMENTED,
        content={"error": {
            "message": f"The '{feature}' feature is not supported by this HugChat-backed API.",
            "type": "not_supported_error",
            "param": None,
            "code": None
        }}
    )

# --- API Endpoints ---

@app.get("/v1/models", response_model=ModelList)
async def list_models():
    if chatbot is None or not available_models_list:
        raise HTTPException(status_code=503, detail="Models list not available. HugChatBot might not be initialized or no models found.")
    
    model_cards = []
    for model_id_str in available_models_list:
        owned_by = "huggingface" # Default
        if "/" in model_id_str:
            # Try to extract owner from "owner/model_name" format
            possible_owner = model_id_str.split('/')[0]
            if possible_owner: # Basic check
                owned_by = possible_owner

        model_cards.append(ModelCard(id=model_id_str, owned_by=owned_by, created=server_start_time))
    
    return ModelList(data=model_cards)

@app.get("/v1/models/{model_id}", response_model=ModelCard)
async def retrieve_model(model_id: str):
    if chatbot is None or not available_models_list:
        raise HTTPException(status_code=503, detail="Model information not available. HugChatBot might not be initialized.")
    
    if model_id in available_models_list:
        owned_by = "huggingface"
        if "/" in model_id:
            possible_owner = model_id.split('/')[0]
            if possible_owner:
                owned_by = possible_owner
        return ModelCard(id=model_id, owned_by=owned_by, created=server_start_time)
    else:
        raise HTTPException(status_code=404, detail=f"Model '{model_id}' not found.")


@app.post("/v1/chat/completions") # response_model removed for StreamingResponse flexibility
async def chat_completions_endpoint(request: ChatCompletionRequest):
    global chatbot, current_llm_model_on_chatbot

    if chatbot is None:
        raise HTTPException(status_code=503, detail="HugChatBot is not available. Check server logs.")
    if not available_models_map:
        raise HTTPException(status_code=503, detail="No LLM models loaded from HugChat.")

    requested_model = request.model
    if requested_model not in available_models_map:
        raise HTTPException(
            status_code=400,
            detail=f"Model '{requested_model}' not found. Available models: {', '.join(available_models_list)}"
        )

    if current_llm_model_on_chatbot != requested_model:
        print(f"Switching model from '{current_llm_model_on_chatbot}' to '{requested_model}'...")
        try:
            model_index = available_models_map[requested_model]
            chatbot.switch_llm(model_index)
            current_llm_model_on_chatbot = requested_model
            print(f"Model switched. Creating new conversation for model: {current_llm_model_on_chatbot}")
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"Failed to switch model: {e}")

    try:
        chatbot.new_conversation(switch_to=True) # Ensure new conversation for this request
        # convo_info = chatbot.get_conversation_info()
        # print(f"New conversation started. Active model: {convo_info.model}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to create new conversation: {e}")

    last_user_message_content = ""
    # OpenAI typically expects a sequence. We'll primarily use the last user message for hugchat.
    # For a more complex setup, one could try to feed prior messages if hugchat supported it explicitly
    # in a single `chat` call beyond its internal memory.
    for msg in reversed(request.messages):
        if msg.role == "user":
            last_user_message_content = msg.content
            break
    
    if not last_user_message_content:
        # Check for system prompt if no user prompt and it's the only message.
        # Though typically OpenAI clients send at least one user message.
        if len(request.messages) == 1 and request.messages[0].role == "system":
             last_user_message_content = request.messages[0].content # Use system as prompt
        else:
            raise HTTPException(status_code=400, detail="No user message found or suitable prompt in the request.")

    prompt = last_user_message_content
    chat_id = f"chatcmpl-{uuid.uuid4().hex}"
    request_time = int(time.time())

    # Handle unsupported parameters (informatively, but hugchat will ignore them)
    if request.temperature is not None and request.temperature != 1.0: # Default OpenAI temp
        print(f"Info: 'temperature' parameter ({request.temperature}) received but may not be supported by HugChat.")
    if request.max_tokens is not None:
        print(f"Info: 'max_tokens' parameter ({request.max_tokens}) received but may not be supported by HugChat.")
    # ... (similar for other params like top_p, n, etc.)

    if request.stream:
        async def stream_generator():
            try:
                first_chunk_data = ChatCompletionChunk(
                    id=chat_id,
                    created=request_time,
                    model=current_llm_model_on_chatbot,
                    choices=[ChatCompletionChunkChoice(delta=DeltaMessage(role="assistant"))]
                )
                yield f"data: {first_chunk_data.model_dump_json(exclude_none=True)}\n\n"

                full_response_text = ""
                # The hugchat stream yields text chunks
                for chunk_text in chatbot.chat(prompt, stream=True):
                    if isinstance(chunk_text, str):
                        full_response_text += chunk_text
                        chunk_data = ChatCompletionChunk(
                            id=chat_id,
                            created=request_time,
                            model=current_llm_model_on_chatbot,
                            choices=[ChatCompletionChunkChoice(delta=DeltaMessage(content=chunk_text))]
                        )
                        yield f"data: {chunk_data.model_dump_json(exclude_none=True)}\n\n"
                    # Add handling for other types if hugchat stream changes
                
                # print(f"Stream complete. Full text for chat {chat_id}: {full_response_text[:100]}...")

                final_chunk_data = ChatCompletionChunk(
                    id=chat_id,
                    created=request_time,
                    model=current_llm_model_on_chatbot,
                    choices=[ChatCompletionChunkChoice(delta=DeltaMessage(), finish_reason="stop")]
                )
                yield f"data: {final_chunk_data.model_dump_json(exclude_none=True)}\n\n"
                yield "data: [DONE]\n\n"
            except Exception as e:
                print(f"Error during streaming for chat {chat_id}: {e}")
                # Attempt to send an error in the stream if possible (before [DONE])
                # This is non-standard for OpenAI, but useful for debugging
                error_content = f"Error during stream: {str(e)}"
                error_delta = DeltaMessage(content=error_content)
                error_choice = ChatCompletionChunkChoice(delta=error_delta, finish_reason="error") # Custom
                error_chunk = ChatCompletionChunk(
                    id=chat_id, created=request_time, model=current_llm_model_on_chatbot, choices=[error_choice]
                )
                try:
                    yield f"data: {error_chunk.model_dump_json(exclude_none=True)}\n\n"
                except Exception: # If stream already broken
                    pass
                yield "data: [DONE]\n\n" # Always end with [DONE]

        return StreamingResponse(stream_generator(), media_type="text/event-stream")
    else: # Non-streaming
        try:
            # Assuming chatbot.chat() with stream=False returns a result object
            # that has wait_until_done() or .text attribute.
            message_result = chatbot.chat(prompt) # hugchat's non-stream returns a Message object

            response_text: str
            if hasattr(message_result, 'wait_until_done'): # If it's a generator-like object
                response_text = message_result.wait_until_done()
            elif hasattr(message_result, 'text'): # If it's a MessageNode or similar
                response_text = message_result.text
            elif isinstance(message_result, str): # Direct string response
                response_text = message_result
            else:
                print(f"Warning: Unexpected response type from chatbot.chat() (non-stream): {type(message_result)}")
                # Attempt to convert to string as a fallback
                try:
                    response_text = str(message_result)
                except:
                    raise ValueError("Could not extract text from HugChat response.")

            # print(f"Non-streamed response for chat {chat_id} / model {current_llm_model_on_chatbot}: {response_text[:100]}...")
            return ChatCompletionResponse(
                id=chat_id,
                created=request_time,
                model=current_llm_model_on_chatbot,
                choices=[
                    ChatCompletionChoice(
                        message=ResponseMessage(role="assistant", content=response_text)
                    )
                ],
                usage=UsageInfo() # Mocked usage
            )
        except Exception as e:
            print(f"Error processing non-streaming chat {chat_id}: {e}")
            raise HTTPException(status_code=500, detail=f"Error processing non-streaming chat: {e}")


# --- Placeholder/Not Implemented Endpoints ---
@app.post("/v1/completions")
async def completions_legacy():
    return not_supported_response("Legacy completions (/v1/completions)")

@app.post("/v1/embeddings")
async def create_embeddings():
    return not_supported_response("Embeddings (/v1/embeddings)")

@app.post("/v1/audio/transcriptions")
async def audio_transcriptions():
    return not_supported_response("Audio transcriptions")

@app.post("/v1/audio/translations")
async def audio_translations():
    return not_supported_response("Audio translations")

@app.post("/v1/images/generations")
async def image_generations():
    # Note: HuggingChat *can* have image generation assistants.
    # A more advanced version could try to map this if a specific assistant ID is known
    # and the request format can be adapted. For now, marking as generally not supported.
    return not_supported_response("Image generations (generic API, specific assistants might work via chat)")

@app.get("/v1/files")
async def list_files_openai(): # Renamed to avoid conflict if you had other /files
    return not_supported_response("File listing/management")

@app.post("/v1/files")
async def upload_file_openai():
    return not_supported_response("File upload")

# ... (add more placeholders for fine-tuning, moderations etc. as needed)

if __name__ == "__main__":
    if not os.path.exists(COOKIE_PATH_DIR):
        try:
            os.makedirs(COOKIE_PATH_DIR)
            print(f"Created directory: {COOKIE_PATH_DIR}")
        except OSError as e:
            print(f"Error creating directory {COOKIE_PATH_DIR}: {e}")
            # Decide if to exit or continue if dir creation fails
            # exit(1)

    print("Starting Uvicorn server...")
    print(f"Credentials: EMAIL={'SET' if HF_EMAIL else 'NOT SET'}, PASSWORD={'SET' if HF_PASSWD else 'NOT SET'}")
    print(f"Cookie Path: {os.path.abspath(COOKIE_PATH_DIR)}")
    uvicorn.run(app, host="0.0.0.0", port=7860)