import os
import gradio as gr
from pymatgen.ext.matproj import MPRester
from crystal_toolkit.components.structure import StructureMoleculeComponent
import json
from crystal_toolkit.core.scene import Scene

# Get the API key from environment variable
MATERIALS_PROJECT_API_KEY = os.getenv('MATERIALS_PROJECT_API_KEY')


def search_materials(query):
    with MPRester(MATERIALS_PROJECT_API_KEY) as mpr:
        results = mpr.summary.search(
            chemsys=query,
            fields=["material_id", "formula_pretty"]
        )
    options = {res.material_id: f"{res.formula_pretty} ({res.material_id})" for res in results}
    if not options:
        return gr.update(choices=[], label="Select Material", value=None)
    return gr.update(choices=list(options.keys()), label="Select Material", value=list(options.keys())[0])


def get_material_data(material_id):
    with MPRester(MATERIALS_PROJECT_API_KEY) as mpr:
        material = mpr.get_structure_by_material_id(material_id)
        summary = mpr.summary.get_data_by_id(material_id)
    return material, summary


def generate_structure_html(scene_json_str):
    import uuid
    div_id = 'structure_' + str(uuid.uuid4())
    html = f'''
    <div id="{div_id}" style="width: 500px; height: 500px;"></div>
    <script src="https://unpkg.com/crystaltoolkit@latest/crystaltoolkit.min.js"></script>
    <script type="application/json" id="{div_id}_data">
    {scene_json_str}
    </script>
    <script>
    const sceneData = JSON.parse(document.getElementById("{div_id}_data").textContent);
    const viewer = new crystalToolkit.ThreeJSViewer("{div_id}", sceneData);
    viewer.render();
    </script>
    '''
    return html


def display_material(material_id):
    if not material_id:
        return "", "Please select a material."

    material, summary = get_material_data(material_id)

    # Create StructureMoleculeComponent
    structure_component = StructureMoleculeComponent(material, id="my_structure")

    # Access the scene json from the component
    scene_json = structure_component._initial_data["scene"]

    # What to do here???
    structure_html = "<p></p>"

    # Extract key properties
    properties = {
        "Material ID": material_id,
        "Formula": summary.formula_pretty,
        "Energy Above Hull (eV/atom)": summary.energy_above_hull,
        "Space Group": summary.symmetry.symbol,
        "Band Gap (eV)": summary.band_gap,
        "Formation Energy (eV/atom)": summary.formation_energy_per_atom,
        "Magnetic Ordering": summary.ordering,
        "Total Magnetization (μB/f.u.)": summary.total_magnetization,
        "Experimentally Observed": summary.is_stable,
        "Crystal System": summary.symmetry.crystal_system,
        "Density (g/cm³)": summary.density,
    }
    # Format properties as HTML
    properties_html = "<table>"
    for key, value in properties.items():
        properties_html += f"<tr><td><strong>{key}</strong></td><td>{value}</td></tr>"
    properties_html += "</table>"

    return structure_html, properties_html


with gr.Blocks() as demo:
    gr.Markdown("## Interactive Crystal Viewer")
    with gr.Row():
        query_input = gr.Textbox(label="Search by Chemical System (e.g., 'Ac-Cd-Ge')", value="Ac-Cd-Ge")
        search_button = gr.Button("Search")
    material_dropdown = gr.Dropdown(label="Select Material", choices=[])
    display_button = gr.Button("Display Material")
    with gr.Row():
        structure_viewer = gr.HTML()
        properties_output = gr.HTML()

    search_button.click(
        fn=search_materials,
        inputs=query_input,
        outputs=material_dropdown
    )

    display_button.click(
        fn=display_material,
        inputs=material_dropdown,
        outputs=[structure_viewer, properties_output]
    )

demo.launch()