import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import re

print("🔄 Chargement du modèle...")

# IGNORER LA QUANTIZATION avec un objet vide
model = AutoModelForCausalLM.from_pretrained(
    "padufour/mistral-accor-emails",
    low_cpu_mem_usage=True,
    device_map=None,  # Pas de device_map auto
    trust_remote_code=True
)

tokenizer = AutoTokenizer.from_pretrained("padufour/mistral-accor-emails")

print("✅ Modèle chargé !")

def generate_email(email_type, client_name, company, objective, 
                   segment, priority, history, style, language):
    
    prompt = f"""### Instruction : Rédige un email du type "{email_type}" pour un client professionnel. L'email doit être complet, sans placeholder ni crochet [], et ne doit PAS mentionner le segment de clientèle dans le contenu.

### Input :
Client Information:
* Name: {client_name}
* Company: {company}
* Objective: {objective}
* Priority: {priority}
* History: {history}
* Style: {style}
* Language: {language}

Important: Ne pas inclure de placeholders [], ne pas mentionner le segment, ne pas mettre les coordonnées du client.

### Réponse :"""
    
    inputs = tokenizer(prompt, return_tensors="pt")
    
    outputs = model.generate(
        **inputs,
        max_new_tokens=800,
        temperature=0.7,
        top_p=0.9,
        top_k=50,
        repetition_penalty=1.1,
        do_sample=True,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id
    )
    
    email = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    if "### Réponse :" in email:
        email = email.split("### Réponse :")[1].strip()
    
    email = re.sub(r'\[.*?\]', '', email)
    email = re.sub(r'\b(MICE|AHBO|leisure group)\b', '', email, flags=re.IGNORECASE)
    email = re.sub(r'at\s+[\w\.-]+@[\w\.-]+\.\w+', '', email)
    
    if "ACCOR Team" not in email:
        email += "\n\nBest regards,\nThe ACCOR Team"
    
    email = re.sub(r'\n\s*\n\s*\n', '\n\n', email)
    return email.strip()

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎯 ACCOR Email Generator\n**Powered by Mistral-7B**")
    
    with gr.Row():
        with gr.Column():
            email_type = gr.Dropdown(
                ["Newsletter", "New offers", "theme", "solution"],
                value="Newsletter", label="Email Type"
            )
            client_name = gr.Textbox(label="Client Name", value="Marie Dupont")
            company = gr.Textbox(label="Company", value="Event Solutions")
            objective = gr.Textbox(label="Objective", value="Promote events", lines=2)
            
            segment = gr.Dropdown(["MICE", "AHBO", "leisure group"], value="MICE", label="Segment")
            priority = gr.Radio(["High", "Medium", "Low"], value="High", label="Priority")
            history = gr.Textbox(label="History", value="New customer")
            style = gr.Textbox(label="Style", value="Professional")
            language = gr.Radio(["ENGLISH", "FRENCH"], value="ENGLISH")
            
            btn = gr.Button("🚀 Generate", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(label="Email", lines=30, show_copy_button=True)
    
    btn.click(generate_email, 
              [email_type, client_name, company, objective, segment, priority, history, style, language],
              output)

demo.launch()