import os
import requests
import gradio as gr

MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY", "")

SYSTEM_PROMPT = (
    "Tu es un expert en physique spécialisé dans les niveaux prépa et licence "
    "(CPGE, L1-L3). Résous les problèmes de manière rigoureuse en détaillant "
    "chaque étape. Structure ta réponse ainsi :\n"
    "**1. Données et inconnues**\n"
    "**2. Modèle physique et hypothèses**\n"
    "**3. Mise en équation**\n"
    "**4. Résolution**\n"
    "**5. Vérification et interprétation**\n\n"
    "Termine toujours par '**Réponse finale :** ...' sur une ligne séparée. "
    "Utilise le LaTeX inline avec $...$ pour les formules."
)

DOMAINES = [
    "⚙️ Mécanique classique",
    "🌡️ Thermodynamique",
    "⚡ Électromagnétisme",
    "🔭 Optique",
    "⚛️ Physique quantique",
    "🚀 Relativité",
]

EXEMPLES = [
    [
        "Un ressort de constante k = 200 N/m est comprimé de x = 5 cm. "
        "Calculer l'énergie potentielle élastique stockée.",
        "⚙️ Mécanique classique", "Prépa (CPGE)",
    ],
    [
        "Une mole de gaz parfait diatomique (gamma = 7/5) subit une compression "
        "adiabatique réversible de V1=10L à V2=2L à partir de T1=300K et P1=1bar. "
        "Calculer T2, P2 et le travail échangé.",
        "🌡️ Thermodynamique", "Prépa (CPGE)",
    ],
    [
        "Un électron est accéléré sous une ddp de U = 50 kV. "
        "Calculer sa longueur d'onde de De Broglie. "
        "Données : me = 9.11e-31 kg, h = 6.626e-34 J.s.",
        "⚛️ Physique quantique", "Licence (L1-L2)",
    ],
    [
        "Un rayon passe de l'eau (n1=1.33) vers le verre (n2=1.5) "
        "avec un angle d'incidence de 30 degrés. "
        "Calculer l'angle de réfraction et vérifier la réflexion totale.",
        "🔭 Optique", "Prépa (CPGE)",
    ],
]

def solve(problem, domaine, niveau):
    if not problem.strip():
        return "⚠️ Entrez un problème de physique."
    if not MISTRAL_API_KEY:
        return "❌ Clé API Mistral manquante (Secret MISTRAL_API_KEY)."
    user_msg = f"Domaine : {domaine}\nNiveau : {niveau}\n\nProblème :\n{problem}"
    try:
        resp = requests.post(
            "https://api.mistral.ai/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {MISTRAL_API_KEY}",
                "Content-Type": "application/json",
            },
            json={
                "model": "mistral-small-latest",
                "messages": [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user",   "content": user_msg},
                ],
                "temperature": 0.2,
                "max_tokens": 1500,
            },
            timeout=45,
        )
        resp.raise_for_status()
        return resp.json()["choices"][0]["message"]["content"].strip()
    except requests.exceptions.Timeout:
        return "⏱️ Timeout — réessaie dans quelques secondes."
    except Exception as e:
        return f"❌ Erreur : {str(e)}"

with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="blue",
        font=[gr.themes.GoogleFont("Inter"), "sans-serif"],
    ),
    title="⚛️ Physics Solver",
) as demo:
    gr.HTML("<h1 style='text-align:center'>⚛️ Physics Solver</h1>")
    gr.HTML(
        "<p style='text-align:center;color:#666'>Résolution de problèmes de physique "
        "niveau <b>Prépa / Licence</b> (CPGE, L1-L3)</p>"
    )

    with gr.Row():
        with gr.Column(scale=2):
            problem_input = gr.Textbox(
                label="Problème",
                placeholder="Décris ton problème ici...",
                lines=6,
            )
            with gr.Row():
                domaine_input = gr.Dropdown(
                    label="Domaine", choices=DOMAINES, value=DOMAINES[0]
                )
                niveau_input = gr.Dropdown(
                    label="Niveau",
                    choices=["Prépa (CPGE)", "Licence (L1-L2)", "Licence avancée (L3)"],
                    value="Prépa (CPGE)",
                )
            solve_btn = gr.Button("▶ Résoudre", variant="primary")

        with gr.Column(scale=3):
            output = gr.Markdown(
                label="Solution",
                value="*La solution apparaîtra ici...*",
            )

    gr.Examples(
        examples=EXEMPLES,
        inputs=[problem_input, domaine_input, niveau_input],
        label="📚 Exemples",
    )

    footer = (
        "Mistral fine-tuné sur SciBench, MMLU, GPQA, OlympiadBench | "
        "[Dataset](https://huggingface.co/datasets/Newton2676/physics-solver-cpge) | "
        "[Modèle](https://huggingface.co/Newton2676/mistral-7b-physics)"
    )
    gr.Markdown(footer)

    solve_btn.click(
        fn=solve,
        inputs=[problem_input, domaine_input, niveau_input],
        outputs=output,
    )
    problem_input.submit(
        fn=solve,
        inputs=[problem_input, domaine_input, niveau_input],
        outputs=output,
    )

if __name__ == "__main__":
    demo.launch()