from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch, gradio as gr

# --- Chargement du modèle ---
MODEL_ID = "Gopu-poss/gopu-agent-2k-fdf"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto"
)

# --- FastAPI ---
app = FastAPI()

class Input(BaseModel):
    input: str
    system_prompt: str | None = None

@app.post("/infer")
def infer(data: Input):
    # Concatène prompt système + input utilisateur
    full_prompt = f"{data.system_prompt}\n{data.input}" if data.system_prompt else data.input
    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.7, top_p=0.9)
    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return {"generated_text": text}

# --- Gradio multimodal ---
# Pipeline de reconnaissance vocale (optionnel)
try:
    asr = pipeline("automatic-speech-recognition", model="openai/whisper-small")
except Exception:
    asr = None

def gopu_chat(system_prompt, prompt, audio=None, image=None, video=None):
    # Si audio fourni → transcrire et remplacer le prompt
    if audio is not None and asr is not None:
        transcription = asr(audio)["text"]
        prompt = transcription

    # Construire le prompt complet
    full_prompt = f"{system_prompt}\n{prompt}" if system_prompt else prompt

    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.7, top_p=0.9)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

demo = gr.Interface(
    fn=gopu_chat,
    inputs=[
        gr.Textbox(label="Prompt système (contexte)", placeholder="Ex: Tu es GopuOS, un agent exclusif..."),
        gr.Textbox(label="Texte utilisateur"),
        gr.Audio(label="Voix (optionnel)", type="filepath"),
        gr.Image(label="Image (optionnel)", type="filepath"),
        gr.Video(label="Vidéo (optionnel)")
    ],
    outputs="text",
    title="GopuOS Agentic Endpoint",
    description="Ajoute un prompt système, parle avec ta voix, ou envoie texte/image/vidéo."
)

# Monter Gradio dans FastAPI
app = gr.mount_gradio_app(app, demo, path="/gradio")