import os
from typing import Optional
from fastapi import FastAPI
from pydantic import BaseModel
from huggingface_hub import login
import requests

API_URL = "https://router.huggingface.co/v1/chat/completions"
headers = {
    "Authorization": f"Bearer {os.environ['HF_TOKEN']}",
}

# Login if token is available (e.g., via Hugging Face Spaces secrets)
if "HF_TOKEN" in os.environ:
    login(token=os.environ["HF_TOKEN"])

app = FastAPI()

class Prompt(BaseModel):
    text: str
    instructions: Optional[str] = None
    output_format: Optional[str] = None

def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()

@app.get("/")
def read_root():
    return {"Hello": "docker embeddings!"}

@app.post("/generate")
async def generate_text(prompt: Prompt):
    messages = []

    # Add system instructions if provided
    if prompt.instructions:
        messages.append({
            "role": "system",
            "content": prompt.instructions
        })

    # Add user message
    user_content = prompt.text

    # If an output format is specified, append formatting instructions
    if prompt.output_format:
        user_content += f"\n\nPlease always format the response as: {prompt.output_format}"

    messages.append({
        "role": "user",
        "content": user_content
    })

    # Call HF API
    response = query({
        "messages": messages,
        "model": "meta-llama/Llama-3.2-3B-Instruct"
    })

    return {"response": response["choices"][0]["message"]}