import js
import os
import ast
import time
import embera
import asyncio
import traceback

import pyscript as ps
import networkx as nx
import minorminer.utils as mmutils 

from pyodide.http import pyfetch

""" ======================================================================= """
""" =========================== Prefecth cache ============================ """

    
import json


async def load_directory():
    rootdir = "assets/cache"
    response = await pyfetch(f"{rootdir}/manifest.json")
    cache = await response.json()
    
    cachedir = 'busclique_' + str(cache['version'])
    os.makedirs(cachedir, exist_ok=True)

    for type in cache['files']:
        os.makedirs(f'{cachedir}/{type}', exist_ok=True)

        for name in cache['files'][type]:
            res = await pyfetch(f"{rootdir}/{cachedir}/{type}/{name}")
            with open(f"{cachedir}/{type}/{name}", "wb") as f:
                f.write(await res.bytes())

# Usage
await load_directory()

""" ======================================================================= """
""" ===================== Encapsulated global data ======================== """
state = embera.AppState()
table = embera.EmbeddingLog("#embedding-table")

""" ======================================================================= """
""" ======================= Source Graph Selector ========================= """

def action_source(event):
    """Updates the source textarea based on dropdown selection."""
    source_sel = ps.document.querySelector("#source-dropdown")
    graph_gen = embera.SOURCE_TYPES[source_sel.value]
    default_layout = embera.DEFAULT_LAYOUT[source_sel.value]
    graph_dim = embera.SOURCE_SIZES[source_sel.value][0]

    graph_name = source_sel.value + str(graph_dim)
    source_graph = graph_gen(*graph_dim)
    source_graph.graph.update({'name':graph_name, 'dim':graph_dim})
    ps.document.querySelector("#source-input").value = source_graph.edges()

    # Clear current size
    spec_sel = ps.document.querySelector("#size-dropdown")
    spec_sel.innerHTML = ""
    
    # Populate new options
    for size in embera.SOURCE_SIZES[source_sel.value]:
        option = ps.document.createElement("option")
        option.value = size
        option.innerText = str(size).replace(',)',')')
        spec_sel.appendChild(option)
    
    # Draw in default mode
    ps.document.querySelector("#layout-dropdown").value = default_layout
    state.S.update_render(graph=source_graph, layout=default_layout)

    # Invalidate embedding
    state.embedding.clear()
    state.T.update_render(emb={})

def action_size(event):
    source_sel = ps.document.querySelector("#source-dropdown")
    size_sel = ps.document.querySelector("#size-dropdown")
    # Delete placeholder entry on first selection
    if not size_sel.children[0].value:
        size_sel.children[0].remove()
    
    graph_gen = embera.SOURCE_TYPES[source_sel.value]
    graph_dim = ast.literal_eval(size_sel.value)
    default_layout = embera.DEFAULT_LAYOUT[source_sel.value]
    graph_name = source_sel.value + size_sel.value

    source_graph = graph_gen(*graph_dim)
    source_graph.graph.update({'name': graph_name,'dim':graph_dim})
    ps.document.querySelector("#source-input").value = source_graph.edges()

    # Draw in default mode
    ps.document.querySelector("#layout-dropdown").value = default_layout
    state.S.update_render(graph=source_graph, layout=default_layout)
    

def action_layout(event):
    layout_sel = ps.document.querySelector("#layout-dropdown")

    # Delete placeholder entry on first selection
    if not layout_sel.children[0].value:
        layout_sel.children[0].remove()

    state.S.update_render(layout=layout_sel.value, emb=state.embedding)   

""" ======================================================================= """
""" ===================== Target Graph Selector =========================== """

def action_target_arch(event):
    """ Updates the 'spec' dropdown for the corresponding 'architecture' """
   
    # Clear current options
    spec_sel = ps.document.querySelector("#target-spec")
    spec_sel.innerHTML = ""
    
    # Populate new options
    arch_val = ps.document.querySelector("#target-arch").value
    for spec in embera.DEVICE_SPECS[arch_val]:
        option = ps.document.createElement("option")
        option.value = spec
        option.innerText = spec
        spec_sel.appendChild(option)
    
    # Invalidate embedding
    state.embedding.clear()

    # Populate textarea with default option
    action_target_spec(None)

def action_target_spec(event):
    """ Update the target textarea based on the specific hardware specs """
    arch_val = ps.document.querySelector("#target-arch").value
    spec_sel = ps.document.querySelector("#target-spec")
    textarea = ps.document.querySelector("#target-input")
    
    target_graph = embera.DEVICE_SPECS[arch_val][spec_sel.value](None)
    textarea.value = list(target_graph.edges())

    # Delete placeholder entry on first selection
    if not spec_sel.children[0].value:
        spec_sel.children[0].remove()

    # Invalidate embedding
    state.embedding.clear()

    # Update
    state.T.update_render(target=target_graph)

""" ======================================================================= """
""" ======================== Rendering Update ============================= """

def update_source(event):
    source_select = js.document.getElementById("source-input")
    layout_select = js.document.getElementById("layout-dropdown")
    source_text = source_select.value.strip() or source_select.placeholder
    try:
        source_val = ast.literal_eval(source_text)
        source_graph = nx.Graph()
        if isinstance(source_val,dict):
            js.console.log("Parsing graph from dictionary")
            key = next(iter(source_val))
            if isinstance(key,tuple):
                for (u, v), w in source_val.items():
                    source_graph.add_edge(u, v, weight=w)
        elif isinstance(source_val,list):
            js.console.log("Parsing graph from edge list")
            source_graph.add_edges_from(source_val)
        source_graph.name = 'Custom'

        state.S.update_render(source_graph, layout_select.value)
        js.document.getElementById("status").innerText = f"Source graph parsed"
    except Exception as e:
        js.console.log(f"Error: {e}")
        js.document.getElementById("status").innerText = f"Error parsing edges"

def update_target(event):
    arch_val = ps.document.querySelector("#target-arch").value
    spec_val = ps.document.querySelector("#target-spec").value
    targetInput = js.document.getElementById("target-input")
    target_text = targetInput.value.strip() or targetInput.placeholder
    
    try:
        target_edges = ast.literal_eval(target_text)
        target_graph = embera.DEVICE_SPECS[arch_val][spec_val](target_edges)
        state.T.update_render(target=target_graph)
        js.document.getElementById("status").innerText = f"Target graph parsed"
    except Exception as e:
        js.console.log(f"Error: {e}")
        js.document.getElementById("status").innerText = f"Error parsing edges"

""" ======================================================================= """
""" ======================== Embedding Action ============================= """
async def run_embedding(event):

    global state
    global table

    method_val = ps.document.querySelector("#emb-method").value
    js.document.getElementById("status").innerText = "Running embedding..."
    await asyncio.sleep(0.05)

    start_time = time.perf_counter()
    try:
        embedding = embera.find_embedding(state.S,state.T,method_val)
        wall_time = time.perf_counter() - start_time    

        if not embedding:
            js.document.getElementById("status").innerText = f"Embedding not found."
            diagnosis = "Empty"
        elif mmutils.verify_embedding(embedding,state.S.edges(),state.T.edges()):
            js.document.getElementById("status").innerText = f"Embedding found."
            diagnosis = "Valid"

    except mmutils.exceptions.EmbeddingError as e:
        js.document.getElementById("status").innerText = repr(e)
        js.console.log(traceback.format_exc())
        diagnosis = type(e).__name__
    except Exception as e:
        wall_time = time.perf_counter() - start_time    
        js.document.getElementById("status").innerText = repr(e)
        js.console.log(traceback.format_exc())
        diagnosis = type(e).__name__
        embedding = {}


    state.S.update_render(emb=embedding)
    state.T.update_render(emb=embedding, source=state.S)
    state.embedding.update_render(embedding)

    results = {"source":state.S.name.replace(',)',')'),
               "edges":state.S.edges(),
               "target":state.T.name,
               "couplers":state.T.edges(),
               "num_nodes":len(state.S),
               "num_edges":len(state.S.edges()),
               "num_qubits":f'{state.embedding.metrics[0]} / {len(state.T)}',
               "num_couplers":len(state.T.edges()), #TODO: interactions
               "method":method_val,
               "time": f'{wall_time:.4f}',
               "embedding":str(embedding),
               "diagnosis": diagnosis,
               "min_chain": state.embedding.metrics[1] if embedding else 0,
               "avg_chain":f"{state.embedding.metrics[0]/len(state.S):.2f}",
               "max_chain":state.embedding.metrics[2]}

    table.addRow(results)


""" ======================================================================= """
""" ======================= Initialization sequence ======================= """
""" ======================================================================= """

js.document.getElementById("status").innerText = "Initializing plots..."

sourceInput = js.document.getElementById("source-input")
source_graph = nx.Graph(ast.literal_eval(sourceInput.placeholder),name='clique(3)',dim=(3,))
state.S = embera.minorplot.SourcePlot(source_graph)

arch_val = js.document.getElementById("target-arch").value
target_graph = embera.DEVICE_SPECS[arch_val]['Unit Cell'](None)
state.T = embera.minorplot.TargetPlot(target_graph) 

state.embedding = embera.minorplot.EmbeddingMap()

js.document.getElementById("status").innerText = "Ready."