import gradio as gr
from huggingface_hub import InferenceClient
import os

HF_TOKEN = os.environ.get("HF_TOKEN")
MODEL_ID  = "Wessym/prompt-refiner-mistral7b"

client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)

SYSTEM_PROMPT = (
    'You are an expert Prompt Engineer. '
    'Your role is to analyze user inputs and either:\n\n'
    '1. REFINE valid prompts into a structured format with these exact sections:\n'
    '   [ROLE]: Who should respond\n'
    '   [CONTEXT]: Background and environment\n'
    '   [TASK]: What exactly needs to be done\n'
    '   [CONSTRAINTS/STYLE]: Style, tone, format constraints\n'
    '   [VARIABLES]: Customizable parameters\n'
    '   [OUTPUT FORMAT]: Expected structure of the output\n\n'
    '   For French inputs use: '
    '[RÔLE], [CONTEXTE], [TÂCHE], [CONTRAINTES/STYLE], [VARIABLES], [FORMAT DE SORTIE]\n\n'
    '2. DETECT and REJECT invalid inputs using the same structure but with:\n'
    '   [ROLE]: Prompt Validator (or [RÔLE]: Validateur de Prompt)\n'
    '   [OUTPUT FORMAT]: INVALID — <TYPE> — <Reason> | Suggestion: <tip>\n\n'
    '   Invalid types: RANDOM_CHARS, NUMBERS_ONLY, PUNCTUATION_ONLY, SOCIAL_GREETING,\n'
    '   TOO_VAGUE, ETHICAL_REFUSAL, REPEATED_CHARS, GENERAL_QUESTION\n\n'
    'Always respond in the same language as the input.'
)

def refine_prompt(user_input: str) -> str:
    if not user_input.strip():
        return "❌ Veuillez entrer un prompt."
    try:
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": f"User input to process: {user_input}"}
        ]
        response = client.chat_completion(
            messages    = messages,
            max_tokens  = 512,
            temperature = 0.3,
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

with gr.Blocks(title="Prompt Refiner") as demo:
    gr.Markdown("# 🚀 Prompt Refiner\nEntre ton prompt et il sera structuré automatiquement.")
    with gr.Row():
        input_box  = gr.Textbox(label="Ton prompt", placeholder="Ex: Crée un script YouTube sur l'IA...", lines=3)
        output_box = gr.Textbox(label="Prompt raffiné", lines=10)
    btn = gr.Button("Raffiner ✨", variant="primary")
    btn.click(fn=refine_prompt, inputs=input_box, outputs=output_box)

demo.launch()