import json
import networkx as nx
# import plotly.graph_objects as go
from pyvis.network import Network
from pathlib import Path


def create_network(data, include_types, type_to_position):
    G = nx.DiGraph()
    for record in data:
        asset_type = record.get("asset_type")
        parent_type = record.get("parent_type")
        child = record.get("external_id")
        parent = record.get("parent_external_id")
        # Filter nodes based on allowed types.
        add_child = include_types is None or asset_type in include_types
        add_parent = include_types is None or parent_type in include_types

        if add_child and child:
            node_attrs = {"type": asset_type, "label": record.get("asset_name")}
            if record.get("uri"):
                node_attrs["uri"] = record.get("uri")
            G.add_node(child, **node_attrs)
        if parent and add_parent:
            node_attrs = {"type": parent_type, "label": record.get("parent_name")}
            if record.get("uri"):
                node_attrs["uri"] = record.get("uri")
            G.add_node(parent, **node_attrs)
        if child and parent and add_child and add_parent:
            G.add_edge(parent, child)

    # Assign 'layer' attribute for each node based on its type.
    for node in G.nodes():
        node_type = G.nodes[node].get('type')
        G.nodes[node]['layer'] = type_to_position.get(node_type, 1)
    return G

# def visualize_network_plotly(data, include_types, type_to_position):
#     """
#     Build and return a Plotly figure of a network graph.

#     Parameters:
#         data (list): JSON array of records.
#         include_types (set): Allowed asset types.
#         type_to_position (dict): Mapping of asset types to numeric layer positions.

#     Returns:
#         go.Figure: The Plotly figure with the network visualization.
#     """
#     # Create network graph

#     G = create_network(data, include_types, type_to_position)
#     # Compute positions using a multipartite layout (vertical alignment)
#     pos = nx.multipartite_layout(G, subset_key="layer", align="vertical")

#     # Build edge coordinates for Plotly
#     edge_x, edge_y = [], []
#     for source, target in G.edges():
#         x0, y0 = pos[source]
#         x1, y1 = pos[target]
#         edge_x.extend([x0, x1, None])
#         edge_y.extend([y0, y1, None])
#     edge_trace = go.Scatter(
#         x=edge_x, y=edge_y,
#         line=dict(width=1, color='#888'),
#         hoverinfo='none',
#         mode='lines'
#     )

#     # Build node coordinates and text lists.
#     node_x, node_y = [], []
#     node_text_trunc, node_text_hover, node_color = [], [], []
#     for node in G.nodes():
#         x, y = pos[node]
#         node_x.append(x)
#         node_y.append(y)
#         full_label = G.nodes[node].get('label', str(node))
#         truncated_label = full_label[0:30]  # truncated version for display
#         # Check for a "uri" attribute (lowercase) and create a clickable link if present.
#         uri = G.nodes[node].get('uri')
#         uri_link = f'<br><a href="{uri}" target="_blank">{uri}</a>' if uri else ""
#         node_text_trunc.append(f"{truncated_label}<br>Type: {G.nodes[node].get('type')}")
#         node_text_hover.append(f"{full_label}<br>Type: {G.nodes[node].get('type')}{uri_link}")
#         node_color.append(G.nodes[node]['layer'])

#     node_trace = go.Scatter(
#         x=node_x, y=node_y,
#         mode='markers+text',
#         textposition="bottom center",
#         hoverinfo='text',
#         marker=dict(
#             showscale=True,
#             colorscale='YlGnBu',
#             reversescale=True,
#             color=node_color,
#             size=15,
#             colorbar=dict(title='Layer')
#         ),
#         text=node_text_trunc,       # Display truncated text on the plot.
#         hovertext=node_text_hover   # Full text (with clickable URL if present) for hover.
#     )

#     # Create the Plotly figure.
#     fig = go.Figure(data=[edge_trace, node_trace],
#                     layout=go.Layout(
#                         title='<br>Integrated Network Graph Visualization',
#                         titlefont_size=16,
#                         dragmode='pan',
#                         showlegend=False,
#                         hovermode='closest',
#                         margin=dict(b=20, l=5, r=5, t=40),
#                         width=1000,
#                         height=800,
#                         xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
#                         yaxis=dict(showgrid=False, zeroline=False, showticklabels=False)
#                     ))
#     return fig

def filter_network(data, include_types, include_fields):
    """
    Filter the data to only include records of the specified types
    and remove specified fields.

    Parameters:
        data (list): List of JSON records.
        include_types (set): Asset types to keep.
        include_fields (set): Fields to remove from records.

    Returns:
        list: Filtered data.
    """
    filtered_data = []
    
    for record in data:
        if record.get("asset_type") in include_types:
            # Create a copy to avoid modifying original data
            filtered_record = {key: value for key, value in record.items() if key in include_fields}
            filtered_data.append(filtered_record)
    
    return filtered_data

def visualize_network_pyvis(data, include_types, type_to_position, output_html = "graph.html"):

    width="100%" 
    height="1200px"
    layer_spacing=2.0
    node_spacing=2.0

    # import tempfile
    # output_html = tempfile.NamedTemporaryFile(delete=False)

    G = create_network(data, include_types, type_to_position)  # Create the network graph
    # 1. Assign hierarchy and determine types
    nx.set_node_attributes(
        G,
        {n: type_to_position.get(attrs["type"], 0) for n, attrs in G.nodes(data=True)},
        name="layer"
    )
    types = sorted({attrs["type"] for _, attrs in G.nodes(data=True)})

    # 2. Layout
    pos = nx.multipartite_layout(G, subset_key="layer", align="horizontal")
    pos = {n: (x * layer_spacing * 1000, y * node_spacing * 1000) for n, (x, y) in pos.items()}

    # 3. Assign colors
    base_colors = ["#e6194B", "#3cb44b", "#ffe119", "#4363d8", "#f58231",
                   "#911eb4", "#46f0f0", "#f032e6", "#bcf60c", "#fabebe"]
    color_map = {t: base_colors[i % len(base_colors)] for i, t in enumerate(types)}

    # 4. Build the graph
    net = Network(width=width, height=height, directed=True)
    net.toggle_physics(False)
    net.toggle_drag_nodes(True)
    net.show_buttons(filter_=["physics", "interaction", "clustering"])

    for node_id, attrs in G.nodes(data=True):
        x, y = pos[node_id]
        t = attrs["type"]
        c = color_map[t]
        net.add_node(node_id, label=attrs.get("label", str(node_id)), title=f"{attrs.get('label')}<br>Type: {t}<br>{attrs.get('uri','')}",
                     x=x, y=y, physics=False, group=t,
                     color={"background": c, "border": c, "highlight": {"background": c, "border": c}, "hover": {"background": c, "border": c}})
    for u, v in G.edges():
        net.add_edge(u, v)

    # 5. Inject clustering + double-click expand
    cluster_js = f"""
<script type="text/javascript">
window.addEventListener('load', () => {{
  const types = {types};
  const colors = {color_map};
  types.forEach(t => {{
    network.cluster({{
      joinCondition: opts => opts.group === t,
      clusterNodeProperties: {{
        id: 'cluster:' + t,
        label: 'Cluster: ' + t,
        shape: 'box',
        color: {{
          background: colors[t],
          border: colors[t],
          highlight: {{background: colors[t], border: colors[t]}},
          hover: {{background: colors[t], border: colors[t]}}
        }}
      }}
    }});
  }});
  network.on('doubleClick', params => {{
    if (params.nodes.length === 1) {{
      const id = params.nodes[0];
      if (network.isCluster(id)) {{
        network.openCluster(id);
      }}
    }}
  }});
}});
</script>
"""

    # 6. Inject a floating HTML legend panel
    legend_html = "<div id='legend' style='position:fixed; top:10px; right:10px; background:white; border:1px solid #ccc; padding:10px; font-family:Arial;'><strong>Legend</strong><br>"
    for t in types:
        c = color_map[t]
        legend_html += f"<div style='display:flex; align-items:center; margin-top:5px;'><span style='width:15px; height:15px; background:{c}; border:1px solid #000; display:inline-block; margin-right:5px;'></span>{t}</div>"
    legend_html += "</div>"
    
    # net.toggle_physics(True)


    # Combine HTML parts
    net.save_graph(output_html)
    with open(output_html, "r+", encoding="utf-8") as f:
        html = f.read().replace("</body>", cluster_js + legend_html + "\n</body>")
        f.seek(0); f.write(html); f.truncate()

    return output_html

def create_lineage_graph(lineage_data, type="pyvis", output_file="graph.html"):

    all_types = {record.get("asset_type") for record in lineage_data if record.get("asset_type")}
    # Exclude unwanted types.
    exclude_types = {"user", "secret", "org_secret", "commit", "actor", "build_machine"}
    types_to_position = {
            "organization": 1,
            "repo": 2,
            "branch": 3,
            "workflow": 4,
            "workflow_run": 5,
            "image": 6,
            "pod": 7,
            "namespace": 8

    }
    include_types = all_types - exclude_types
    # Generate and display the network visualization.
    if type == "plotly":
        # fig = visualize_network_plotly(lineage_data, include_types, types_to_position)
        return None
    else:
        output_file = Path(output_file)
        output_file.touch(exist_ok=True)
        output_file = output_file.resolve()
        output_file = str(output_file)
        file = visualize_network_pyvis(lineage_data, include_types, types_to_position, output_html=output_file)
        return file

if __name__ == "__main__":
    # Open the JSON file and load data.
    filepath = "demo-data/astro-lineage.json"
    with open(filepath, "r") as f:
        data = json.load(f)
    

    filename = create_lineage_graph(data, output_file="demo-data/lineage_graph.html")
    print(f"Network graph saved to {filename}")
    print(f"Open by browsing file:///{filename}")