"""
Molecule Maestro – Category Shuffle Edition (Markdown‑Table Output)
------------------------------------------------------------------
* Six category boxes (🔴 Acids, 🔵 Bases, …).
* **Reshuffle** keeps 50 % of reactants.
* One free‑form **Conditions** textbox.
* **NEW:** GPT now returns a ready‑made **Markdown table** – no JSON parsing required.

```bash
pip install "gradio>=4" "openai>=1.0"
export OPENAI_API_KEY="sk‑…"
python molecule_maestro_gradio.py
```
"""

from __future__ import annotations
import random
from typing import List, Dict
from openai import OpenAI
import gradio as gr

client = OpenAI()

CATEGORIES = [
    ("acid", "🔴 Acids"),
    ("base", "🔵 Bases"),
    ("solvent", "🟢 Solvents"),
    ("gas", "🟡 Gases"),
    ("oxidizer", "🟣 Oxidizers"),
    ("other", "⚪ Others"),
]

SYSTEM_PROMPT_PALETTE = (
    "You are a helpful chemistry assistant. Provide a JSON array named 'palette' with exactly 12 reactants commonly used in school labs. "
    "Each item must have fields: name (string) and category (acid, base, solvent, gas, oxidizer, or other)."
)

SYSTEM_PROMPT_REACTION = (
    "You are a chemistry mentor. The user will supply 1‑5 reactants and an optional conditions string. "
    "Respond ONLY with a Markdown table **including the header row**: \n"
    "| Reactants | Conditions | Reaction Equation | Observations |\n"
    "|---|---|---|---|\n"
    "Each subsequent row lists **one** plausible, single‑step, school‑level reaction (max 5 rows total). "
    "Observations must indicate exothermic or endothermic and an approximate ΔH in kJ, plus a ≤25‑word note. "
    "If no simple reaction is feasible, output the header row followed by a row stating 'No simple reaction found' in the Reaction Equation column and leave other cells blank."
)


def chat(prompt: str, system: str, temp: float = 0.4) -> str:
    return client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role": "system", "content": system}, {"role": "user", "content": prompt}],
        temperature=temp,
    ).choices[0].message.content.strip()


# ---------- Palette helpers ----------

def gen_palette() -> List[Dict[str, str]]:
    try:
        import json
        raw = chat("Give me the palette", SYSTEM_PROMPT_PALETTE)
        return json.loads(raw)["palette"]
    except Exception:
        return [
            {"name": "HCl", "category": "acid"}, {"name": "NaOH", "category": "base"}, {"name": "Ethanol", "category": "solvent"},
            {"name": "Acetic Acid", "category": "acid"}, {"name": "NH3", "category": "base"}, {"name": "H2O", "category": "solvent"},
            {"name": "H2", "category": "gas"}, {"name": "O2", "category": "gas"}, {"name": "KMnO4", "category": "oxidizer"},
            {"name": "H2SO4", "category": "acid"}, {"name": "NaCl", "category": "other"}, {"name": "CuSO4", "category": "other"},
        ]

palette: List[Dict[str, str]] = gen_palette()


def labels_by_cat(cat: str) -> List[str]:
    return [item["name"] for item in palette if item["category"] == cat]


# ---------- Reaction evaluation ----------

def evaluate(reactants: List[str], cond: str) -> str:
    if not reactants:
        return "⚠️ Select at least one reactant."

    prompt = f"Reactants: {', '.join(reactants)}. Conditions: {cond or 'none'}."
    table = chat(prompt, SYSTEM_PROMPT_REACTION, 0.3)
    # Basic sanity check: ensure header present
    if "| Reactants |" not in table:
        return f"Parse err 🤖\n```\n{table}\n```"
    return table


# ---------- Gradio UI ----------
with gr.Blocks() as demo:
    gr.Markdown("# 🧪 Molecule Maestro – Category Shuffle Edition")

    with gr.Row():
        cat_groups = {}
        for cat, title in CATEGORIES:
            with gr.Column():
                cat_groups[cat] = gr.CheckboxGroup(choices=labels_by_cat(cat), label=title)
        with gr.Column():
            shuffle_btn = gr.Button("🔀 Reshuffle (50 % new)")
            conditions_tb = gr.Textbox(label="Conditions (optional)")
            run_btn = gr.Button("▶︎ Generate Reactions")

    result_md = gr.Markdown("—")

    # ---------- Callbacks ----------
    def _shuffle():
        global palette
        keep = random.sample(palette, len(palette)//2)
        kept_names = {i['name'] for i in keep}
        fresh: List[Dict[str, str]] = []
        attempts = 0
        while len(fresh) < len(palette)//2 and attempts < 5:
            for item in gen_palette():
                if item['name'] not in kept_names and item['name'] not in {f['name'] for f in fresh}:
                    fresh.append(item)
                    if len(fresh) == len(palette)//2:
                        break
            attempts += 1
        palette[:] = keep + fresh
        random.shuffle(palette)
        updates = [gr.update(choices=labels_by_cat(cat), value=[]) for cat, _ in CATEGORIES]
        updates.append("—")
        return updates

    def _run(*inputs):
        *checkbox_lists, cond = inputs
        chosen: List[str] = []
        for lst in checkbox_lists:
            chosen.extend(lst)
        return evaluate(chosen, cond)

    shuffle_btn.click(_shuffle, None, [cat_groups[c] for c, _ in CATEGORIES] + [result_md])
    run_btn.click(_run, [cat_groups[c] for c, _ in CATEGORIES] + [conditions_tb], result_md)

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