import molfoundry as mf
import streamlit as st
import pandas as pd
import requests
from rdkit import Chem
from rdkit.Chem import Draw

import json
from pathlib import Path
import re


def parse_rules():
    rules = {}
    folder = Path("../2026-OverlayGraphs/learning_exports/aams/ogs")
    MAX = 30
    count = 0
    for p in folder.glob("*.gml"):
        if count>=MAX: break
        count+=1      # only .gml files in this folder
        text = p.read_text()
        rule_gml = text
        m = re.search(r"https?://[^\s<>\"')\]]+", text)
        url = m.group(0).rstrip(".,;:!?") if m else None
        m = re.search(r"/entry/(\d+)", url)
        mcsa_id = m.group(1) if m else None
        m = re.search(r"ruleID \"(.*)\"", text)
        rule_id = m.group(1) if m else None
        rules[int(mcsa_id)] = {'mcsa_id': mcsa_id, 'rule_id': rule_id, 'url': url, 'gml_string': rule_gml}
    return rules

def parse_smiles():
    smiles = []
    mols = []
    with open("./output/mcsa-smiles.json", 'r') as f:
        data = json.load(f)

    for entry in data:
        #g = mf.smiles(data[entry]['smiles'], entry, add=False)
        smiles.append(data[entry]['smiles'])
        mols.append(Chem.MolFromSmiles(data[entry]['smiles']))
    return (smiles, mols)

def get_compounds_mcsa(mcsa_id):
    BASE = "https://www.ebi.ac.uk/thornton-srv/m-csa/api"
    r = requests.get(f"{BASE}/entries/{mcsa_id}/", params={"format": "json"}, timeout=60)
    r.raise_for_status()
    entry = r.json()
    reaction = entry.get("reaction")
    compounds = reaction["compounds"]

    reactants=[]
    products = []
    for comp in compounds:
        cnt = comp['count']
        comp_type = comp['type']
        name = comp['name']
        mol_url = comp['mol_file']
        if "reactant" in comp_type:
            reactants.append(comp)
        else:
            products.append(comp)

    return (reactants, products)

def draw_molfile(url):
    url = f"https://{url}"
    r = requests.get(url, timeout=30)
    r.raise_for_status()
    molblock = r.text

    mol = Chem.MolFromMolBlock(molblock)
    if mol is None:
        raise ValueError("RDKit could not parse the MOL file")

    return Draw.MolToImage(mol, size=(200, 150))

def main():
    rules = parse_rules()
    mcsa_ids = sorted([int(rules[r]['mcsa_id']) for r in rules])

    st.set_page_config(layout="wide")

    df = pd.DataFrame(mcsa_ids, columns=["mcsa_id"])

    if "selected" not in st.session_state:
        st.session_state.selected = None

    if "op_selected" not in st.session_state:
        st.session_state.op_selected = None

    def select(mcsa_id):
        st.session_state.selected = mcsa_id

    def back():
        st.session_state.selected = None
        st.session_state.op_selected = None


    # ---------- full-screen view ----------
    if st.session_state.selected is not None:
        sel = st.session_state.selected
        st.button("← Back to grid", on_click=back)
        st.subheader("Reaction Comparison")
        st.write(f"**{rules[sel]['rule_id']}**")
        st.markdown(f"[{rules[sel]['url']}]({rules[sel]['url']})")
        (reactants, products) = get_compounds_mcsa(sel)

        # split into reactants and products
        left, right = st.columns(2)
        with left:
            st.subheader("Reactants")
            for r in reactants:
                box = st.container()
                with box:
                    st.image(draw_molfile(r['mol_file']))
                    st.write(f"**Count:** {r['count']}, **Name:** {r['name']}, **Chebi:** {r['chebi_id']}")
        with right:
            st.subheader("Products")
            for p in products:
                box = st.container()
                with box:
                    st.image(draw_molfile(p['mol_file']))
                    st.write(f"**Count:** {p['count']}, **Name:** {p['name']}, **Chebi:** {p['chebi_id']}")
        
        mfbox = st.container()
        with mfbox:
            st.subheader("OG products")
            left, right = st.columns([1,10])
            with left:
                options = [Path(p).name for p in Path(f"output/mf_products/{sel}").glob("*")]
                for op in options:
                    if left.button(op, key=f"btn_{op}", use_container_width=True):
                        st.session_state.op_selected = op
            with right:
                if st.session_state.op_selected:
                    op = st.session_state.op_selected
                    if "match" in op:
                        with open(f"output/mf_products/{sel}/match.txt", 'r') as f:
                            data = f.read()
                        st.code(data)
                    elif "reaction" in op:
                        with open(f"output/mf_products/{sel}/reaction.svg", 'r') as f:
                            svg = f.read()
                        st.image(svg, width=1000)
                    else:
                        folder = Path(f"output/mf_products/{sel}/{op}")
                        for p in folder.glob("*.svg"):
                            with open(p, 'r') as f:
                                svg = f.read()
                            st.image(svg, width=300)
                else:
                    st.write("Choose a derivation to see products.")
        st.stop()          # nothing below this line runs in this view

    # ---------- grid view ----------
    st.title("M-CSA IDs")
    N_COLS = 9
    ids = df["mcsa_id"].tolist()
    for start in range(0, len(ids), N_COLS):
            cols = st.columns(N_COLS)
            for col, cid in zip(cols, ids[start:start + N_COLS]):
                if not Path(f"output/mf_products/{cid}/match.txt").exists(): continue
                with open(f"output/mf_products/{cid}/match.txt", 'r') as f:
                    green = f.read()
                if "TRUE" in green:
                    if col.button(f":green[{cid}]", key=f"btn_{cid}", use_container_width=True):
                        st.session_state.selected = cid
                else:
                    if col.button(f":red[{cid}]", key=f"btn_{cid}", use_container_width=True):
                        st.session_state.selected = cid

main()