import base64
import dash
from dash import html, dcc, dash_table
import dash_cytoscape as cyto
from dash.dependencies import Input, Output, State
import pandas as pd

from utility import parse_functions_from_files, get_reachable_functions, get_backtrace_functions, build_nodes_and_edges
from introduction import get_intro_markdown

app = dash.Dash(__name__, suppress_callback_exceptions=True)
server = app.server


app.layout = html.Div([
    html.H2("Function Dependency Visualizer (AST-Based)", style={
        "textAlign": "center",
        "marginTop": "20px",
        "fontSize": "28px",
        "color": "#333"
    }),
    dcc.Store(id="parsed-functions-store"),

    html.Div([
        html.Div([
            html.Label("Upload Files"),
            dcc.Upload(
                id="upload",
                children=html.Div("📤 Drag and drop or click to upload Python files", style={
                    "border": "2px dashed #ccc",
                    "padding": "10px",
                    "textAlign": "center",
                    "cursor": "pointer",
                    "color": "#555",
                    "fontSize": "14px"
                }),
                multiple=True
            )
        ], style={"width": "24%", "display": "inline-block", "verticalAlign": "top", "marginRight": "1%"}),

        html.Div([
            html.Label("Main Function"),
            dcc.Dropdown(id="main-function", style={"width": "100%"})
        ], style={"width": "24%", "display": "inline-block", "verticalAlign": "top", "marginRight": "1%"}),

        html.Div([
            html.Label("Max Depth"),
            dcc.Dropdown(
                id="max-depth",
                options=[{"label": str(i), "value": i} for i in range(1, 11)],
                value=10,
                clearable=False,
                style={"width": "100%"}
            )
        ], style={"width": "24%", "display": "inline-block", "verticalAlign": "top", "marginRight": "1%"}),

        html.Div([
            html.Label("Options"),
            dcc.Checklist(
                id="backtrace-toggle",
                options=[{"label": "Backtrace Mode", "value": "backtrace"}],
                value=[],
                labelStyle={"display": "block", "fontWeight": "normal"}
            ),
            dcc.Checklist(
                id="function-detail-toggle",
                options=[{"label": "Show Function-Level Detail in Summaray tab", "value": "show"}],
                value=[],
                labelStyle={"display": "block", "fontWeight": "normal"},
                style={"marginTop": "5px"}
            )
        ], style={"width": "24%", "display": "inline-block", "verticalAlign": "top"})
    ], style={
        "backgroundColor": "#f9f9f9",
        "padding": "15px",
        "borderRadius": "10px",
        "margin": "20px auto",
        "width": "95%",
        "boxShadow": "0 2px 6px rgba(0,0,0,0.1)"
    }),


    dcc.Tabs(id="tab-selector", value="intro", children=[
        dcc.Tab(label="📘 Introduction", value="intro"),
        dcc.Tab(label="📊 Graph Explorer", value="graph"),
        dcc.Tab(label="📁 File Summary", value="summary")
    ], style={"margin": "0 20px"}),

    html.Div(html.Div(id="intro-tab", style={"margin": "0 20px"}), id="intro-tab-container"),
    html.Div(id="graph-tab-container", style={"margin": "0 20px"}),
    html.Div(id="summary-tab-container", style={"margin": "0 20px"}),


    html.Footer("Built by Tomas Larsson • MIT Licensed", style={
        "textAlign": "center",
        "fontSize": "14px",
        "marginTop": "40px",
        "marginBottom": "20px",
        "color": "#888"
    })
])


@app.callback(
    Output("main-function", "options"),
    Output("main-function", "value"),
    Output("parsed-functions-store", "data"),
    Input("upload", "contents"),
    State("upload", "filename")
)
def handle_upload(contents, filenames):
    print("=== handle_upload triggered ===")
    if not contents:
        return [], None, {}
    uploaded_files = {}
    for content, name in zip(contents, filenames):
        _, content_string = content.split(",")
        uploaded_files[name] = base64.b64decode(content_string).decode("utf-8")
    parsed = parse_functions_from_files(uploaded_files)


#    options = [{"label": fn, "value": fn} for fn in parsed if parsed[fn]["calls"]]
    options = [{"label": fn, "value": fn} for fn in parsed]
    return options, options[0]["value"] if options else None, parsed

@app.callback(
    Output("intro-tab-container", "style"),
    Output("graph-tab-container", "style"),
    Output("summary-tab-container", "style"),
    Input("tab-selector", "value")
)
def toggle_tabs(tab):
    return (
        {"display": "block"} if tab == "intro" else {"display": "none"},
        {"display": "block"} if tab == "graph" else {"display": "none"},
        {"display": "block"} if tab == "summary" else {"display": "none"}
    )

@app.callback(Output("intro-tab", "children"), Input("tab-selector", "value"))
def show_intro(_):
    return get_intro_markdown()

from collections import deque, defaultdict

@app.callback(
    Output("graph-tab-container", "children"),
    Input("main-function", "value"),
    Input("backtrace-toggle", "value"),
    Input("max-depth", "value"),
    State("parsed-functions-store", "data")
)
def update_graph(main_func, backtrace_mode, max_depth, parsed):
    if not parsed or not main_func:
        return html.Div("Upload files and select a main function.")

    graph = {k: v["calls"] for k, v in parsed.items()}
    reachable = get_backtrace_functions(main_func, graph) if "backtrace" in backtrace_mode else get_reachable_functions(main_func, graph)

    nodes, edges = build_nodes_and_edges(parsed, main_func, reachable, reverse="backtrace" in backtrace_mode, max_depth=max_depth)


    return cyto.Cytoscape(
        id="cytoscape-graph",
        layout={"name": "preset"},
        style={"width": "100%", "height": "600px"},
        elements=nodes + edges,
        stylesheet=[
            {"selector": "node", "style": {
                "label": "data(label)",
                "text-wrap": "wrap",
                "text-valign": "bottom",
                "text-halign": "center"
            }},
            {"selector": ".main", "style": {
                "background-color": "blue",
                "line-color": "blue",
                "color": "black",
                "label": "data(label)",
                "text-wrap": "wrap",
                "text-valign": "bottom",
                "text-halign": "center"
            }},
            {"selector": "edge", "style": {
                "label": "data(label)",
                "curve-style": "bezier",
                "target-arrow-shape": "triangle",
                "target-arrow-color": "#888",
                "line-color": "#888",
                "arrow-scale": 2,
                "font-size": "14px",
                "text-margin-y": -10,
                "text-margin-x": 10,
            }}
        ],
        userZoomingEnabled=True,
        userPanningEnabled=True,
        minZoom=0.2,
        maxZoom=2,
        wheelSensitivity=0.1
    )


@app.callback(
    Output("summary-tab-container", "children"),
    Input("parsed-functions-store", "data"),
    Input("function-detail-toggle", "value")
)
def update_summary(parsed, toggle):
    if not parsed:
        return html.Div("No summary available.")
    reverse_calls = {k: [] for k in parsed}
    for caller, meta in parsed.items():
        for callee in meta["calls"]:
            reverse_calls[callee].append(caller)

    if "show" in toggle:
        df = pd.DataFrame([{
            "Function": fn,
            "File": parsed[fn]["filename"],
            "Arguments": ", ".join(parsed[fn]["args"]),
            "Returns": ", ".join(parsed[fn]["returns"]),
            "Reads State": ", ".join(parsed[fn]["reads_state"]),
            "Writes State": ", ".join(parsed[fn]["writes_state"]),
            "Called By": len(reverse_calls[fn])
        } for fn in parsed])
    else:
        file_summary = {}
        for func, meta in parsed.items():
            fname = meta["filename"]
            file_summary.setdefault(fname, {"Total": 0, "Calls Others": 0, "Called By Others": 0, "Unused": 0})
            file_summary[fname]["Total"] += 1
            if meta["calls"]: file_summary[fname]["Calls Others"] += 1
            if reverse_calls[func]: file_summary[fname]["Called By Others"] += 1
            if not meta["calls"] and not reverse_calls[func]: file_summary[fname]["Unused"] += 1
        df = pd.DataFrame([{"File": f, **stats} for f, stats in file_summary.items()])

    return dash_table.DataTable(
        data=df.to_dict("records"),
        columns=[{"name": c, "id": c} for c in df.columns],
        style_table={"overflowX": "auto"},
        style_cell={"whiteSpace": "normal", "textAlign": "left", "padding": "5px", "maxWidth": 300}
    )

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)