import gradio as gr
import pandas as pd
import time

# --- 模拟后端功能 (来自第二个版本) ---
def smiles_to_structure(smiles):
    """Dummy function"""
    time.sleep(1)
    return f"Structure Generated from: {smiles}"

def fragment_molecule(smiles):
    """Dummy function"""
    time.sleep(2)
    return "Fragment1", "Fragment2", "1-2"

def generate_analogs(main_cls, minor_cls, number, delta_value):
    """Dummy function"""
    time.sleep(3)
    return [
        {"SMILE": "c1cccc1", "MolWt": 100, "TPSA": 20, "SLogP": 1, "SA": 30, "QED": 0.8},
        {"SMILE": "c1ccccc1", "MolWt": 105, "TPSA": 25, "SLogP": 1.2, "SA": 32, "QED": 0.9},
    ]

def update_output_table(data):
    df = pd.DataFrame(data)
    return df

# --- Gradio 界面 ---
KETCHER_HTML = r'''
<iframe id="ifKetcher" src="/ketcher/index.html" width="100%" height="600px" style="border: 1px solid #ccc;"></iframe>

<script>
console.log("[Front-end] Ketcher-Gradio integration script loaded.");

let ketcher = null;
let lastSmiles = '';

function findSmilesInput() {
    const inputContainer = document.getElementById('combined_smiles_input');
    if (!inputContainer) {
        console.warn("[Front-end] combined_smiles_input element not found.");
        return null;
    }
    const input = inputContainer.querySelector('input[type="text"]');
    return input;
}

function updateGradioInput(smiles) {
    const input = findSmilesInput();
    if (input && input.value !== smiles) {
        input.value = smiles;
        input.dispatchEvent(new Event('input', { bubbles: true }));
        console.log("[Front-end] Updated Gradio input with SMILES:", smiles);
    }
}

async function handleKetcherChange() {
    console.log("[Front-end] handleKetcherChange called, retrieving SMILES...");
    try {
        const smiles = await ketcher.getSmiles({ arom: false });
        console.log("[Front-end] SMILES retrieved from Ketcher:", smiles);
        if (smiles !== lastSmiles) {
            lastSmiles = smiles;
            updateGradioInput(smiles);
        }
    } catch (err) {
        console.error("[Front-end] Error getting SMILES from Ketcher:", err);
    }
}

function initKetcher() {
    console.log("[Front-end] initKetcher started.");
    const iframe = document.getElementById('ifKetcher');
    if (!iframe) {
        console.error("[Front-end] iframe not found.");
        setTimeout(initKetcher, 500);
        return;
    }

    const ketcherWindow = iframe.contentWindow;
    if (!ketcherWindow || !ketcherWindow.ketcher) {
        console.log("[Front-end] ketcher not yet available in iframe, retrying...");
        setTimeout(initKetcher, 500);
        return;
    }

    ketcher = ketcherWindow.ketcher;
    console.log("[Front-end] Ketcher instance acquired:", ketcher);

    ketcher.setMolecule('C').then(() => {
        console.log("[Front-end] Initial molecule set to 'C'.");
    });

    const editor = ketcher.editor;
    if (editor && typeof editor.subscribe === 'function') {
        console.log("[Front-end] Using editor.subscribe('change', ...)");
        editor.subscribe('change', handleKetcherChange);
    } else {
        console.error("[Front-end] No suitable event binding found.");
    }
}

document.getElementById('ifKetcher').addEventListener('load', () => {
    console.log("[Front-end] iframe loaded. Initializing Ketcher in 1s...");
    setTimeout(initKetcher, 1000);
});
</script>
'''

def create_combined_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# Fragment Optimization Tools with Ketcher")

        with gr.Row():
            with gr.Column(scale=2):
                gr.HTML(KETCHER_HTML)  # 嵌入 Ketcher，仅加载一次

            with gr.Column(scale=1):
                gr.Markdown("### Input SMILES (From Ketcher)")
                combined_smiles_input = gr.Textbox(
                    label="",
                    value="C",
                    placeholder="SMILES from Ketcher will appear here",
                    elem_id="combined_smiles_input"
                )
                get_ketcher_smiles_btn = gr.Button("Get SMILES from Ketcher")
                fragment_btn = gr.Button("Fragmentize Molecule")

                constant_frag_input = gr.Textbox(label="Constant Fragment", placeholder="SMILES of constant fragment")
                variable_frag_input = gr.Textbox(label="Variable Fragment", placeholder="SMILES of variable fragment")
                attach_order_input = gr.Textbox(label="Attachment Order", placeholder="Attachment Order of SMILES")

                main_cls_dropdown = gr.Dropdown(label="Main Cls", choices=["None", "Cl", "Br"])
                minor_cls_dropdown = gr.Dropdown(label="Minor Cls", choices=["None", "Cl", "Br"])
                number_input = gr.Number(label="Number", value=5, step=1)
                delta_value_slider = gr.Slider(minimum=0, maximum=10, step=1, label="Delta Value", interactive=True)
                generate_analogs_btn = gr.Button("Generate")

        output_table = gr.Dataframe(headers=["SMILE", "MolWt", "TPSA", "SLogP", "SA", "QED"])

        # --- 事件处理 ---
        get_ketcher_smiles_btn.click(
            fn=None,
            inputs=None,
            outputs=combined_smiles_input,
            js="async () => { const iframe = document.getElementById('ifKetcher'); if(iframe && iframe.contentWindow && iframe.contentWindow.ketcher) { const smiles = await iframe.contentWindow.ketcher.getSmiles(); return smiles; } else { console.error('Ketcher not ready'); return ''; } }"
        )

        fragment_btn.click(fragment_molecule, inputs=combined_smiles_input, outputs=[constant_frag_input, variable_frag_input, attach_order_input])

        def update_table_with_analogs(main_cls, minor_cls, number, delta_value):
            analogs_data = generate_analogs(main_cls, minor_cls, number, delta_value)
            return update_output_table(analogs_data)

        generate_analogs_btn.click(update_table_with_analogs,
                                   inputs=[main_cls_dropdown, minor_cls_dropdown, number_input, delta_value_slider],
                                   outputs=output_table)

    return demo


# 创建 Gradio 界面并启动
demo = create_combined_interface()
demo.launch()