import os
import tkinter as tk
from pathlib import Path

import matplotlib
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
from graphviz import Source

import cfg_build
import lib.console as cnsl
import syntax
import triplayacc as yacc
from cfa.analysis_dot import analysis_to_dot, run_all_analyses
from cfg.CFG import CFG
from vistram.tram import *
from vistram.vistram import MachineUI

matplotlib.use("TkAgg")

# Assembles the AST into TRAM code
def assemble(ast):
    code = ast.code({}, 0)
    return code + [halt()]

def make_cfg(ast):
    return CFG(ast)

# Renders a diagram of the AST
def render_diagram(dot_string: str):
    # Set DPI for PNG
    os.environ["GV_FILE_DPI"] = "300"

    src = Source(dot_string, format="png", engine="dot")
    png_path = src.render(cleanup=True)

    img = mpimg.imread(png_path)

    fig = plt.figure(figsize=(12, 12))
    fig.canvas.manager.window.wm_title("TRIPLA AST Viewer")

    plt.imshow(img)
    plt.axis("off")
    plt.show()

# Pretty prints the AST
def pretty_print(node, indent=0):
    prefix = " " * indent
    print(f"{prefix}{type(node).__name__}:")

    for key, value in node.__dict__.items():
        if isinstance(value, syntax.EXPRESSION):
            pretty_print(value, indent + 4)

        elif isinstance(value, list):
            print(f"{prefix}  {key}: [")
            for element in value:
                if isinstance(element, syntax.EXPRESSION):
                    pretty_print(element, indent + 4)
                else:
                    print(" " * (indent + 4) + str(element))
            print(f"{prefix}  ]")

        else:
            print(f"{prefix}  {key}: {value}")


def print_analysis_reports(cfg, analyses: dict, ru_edges: dict[int, list[int]]):
    """Print compact Live Variables and Reached Uses reports to console."""
    lv = analyses["lv"]
    ru = analyses["ru"]
    _ = ru
    node_by_id = {n.id: n for n in cfg.nodes()}

    def node_text(nid: int) -> str:
        node = node_by_id.get(nid)
        if node is None:
            return "<unknown>"
        lbl = node.dot_label()
        return str(lbl) if lbl is not None else "<no-label>"

    print("\nLive Variables Report")
    print("---------------------")
    node_ids = sorted(set(lv.in_sets.keys()) | set(lv.out_sets.keys()))
    for nid in node_ids:
        in_set = sorted(lv.in_sets.get(nid, set()))
        out_set = sorted(lv.out_sets.get(nid, set()))
        if not in_set and not out_set:
            continue
        print(f"n{nid} [{node_text(nid)}]: In={in_set} Out={out_set}")

    print("\nReached Uses Report")
    print("-------------------")
    has_ru = False
    for def_id in sorted(ru_edges):
        uses = sorted(set(ru_edges[def_id]))
        if not uses:
            continue
        has_ru = True
        use_desc = [f"n{uid} [{node_text(uid)}]" for uid in uses]
        print(f"n{def_id} [{node_text(def_id)}] -> {use_desc}")
    if not has_ru:
        print("(no reached uses)")


if __name__ == "__main__":
    print("\nTRIPLA Parser and TRIPLA to TRAM Compiler")

    while True:
        mode = cnsl.prompt_choice(
            "\nSelect action:",
            ["Parse .tripla", "Compile .tripla", "CFG for .tripla", "Analyze .tripla", "Exit"],
        )

        if mode == 4:
            print("\nBye Bye.")
            break

        base = Path(__file__).parent / "triplaprograms"
        programs = sorted([f for f in base.glob("*.tripla")])
        if not programs:
            print("\nNo .tripla files found.")
            continue

        idx = cnsl.prompt_choice("\nSelect TRIPLA program:", [p.name for p in programs])
        path = programs[idx]

        source = path.read_text()
        ast = yacc.parser.parse(source)

        if mode == 0:
            # Pretty print
            if cnsl.prompt_confirmation("\nPretty-print AST?"):
                print("")
                pretty_print(ast)

            # Export DOT
            dot_str = ast.to_dot()
            if cnsl.prompt_confirmation("Export AST as .dot file?"):
                default = f"{path.stem}.dot"
                filename = input(f"Filename [{default}]: ").strip()
                if not filename:
                    filename = default
                try:
                    base_dir = Path(__file__).parent
                    out_path = base_dir / filename
                    with open(out_path, "w") as f:
                        f.write(dot_str)
                    print(f"Saved DOT file as: {out_path}")
                except Exception as e:
                    print(f"Could not save DOT file: {e}")

            # Display AST diagram
            if cnsl.prompt_confirmation("Display AST diagram?"):
                render_diagram(dot_str)
                print("Rendered AST diagram.")

        elif mode == 1:
            tram_code = assemble(ast)

            # Print TRAM code
            if cnsl.prompt_confirmation("\nPrint TRAM code to console?"):
                print("\nGenerated TRAM code:\n")
                for instr in tram_code:
                    print(instr.toString())

            # Save TRAM code
            if cnsl.prompt_confirmation("Save TRAM code as .tram file?"):
                base_dir = Path(__file__).parent / "tramcodes"
                base_dir.mkdir(exist_ok=True)

                default = f"{path.stem}.tram"
                filename = input(f"Filename [{default}]: ").strip()
                if not filename:
                    filename = default

                out_path = base_dir / filename
                with open(out_path, "w") as f:
                    for instr in tram_code:
                        f.write(instr.toString() + "\n")

                print(f"Saved TRAM code to: {out_path}")

            # Display TRAM code in Visual TRAM UI
            if cnsl.prompt_confirmation("Display TRAM code in Visual TRAM UI?"):
                root = tk.Tk()
                ui = MachineUI(root)
                ui.machine.initial_program = tram_code
                ui.machine.reset()
                root.mainloop()

        elif mode == 2:
            cfg = make_cfg(ast)

            dot_str = cfg.to_dot()
            if cnsl.prompt_confirmation("\nExport CFG as .dot file?"):
                default = f"{path.stem}_cfg.dot"
                filename = input(f"Filename [{default}]: ").strip()
                if not filename:
                    filename = default

                out_path = Path(__file__).parent / 'cfgdots' / filename
                out_path.parent.mkdir(exist_ok=True)
                with open(out_path, "w") as f:
                    f.write(dot_str)

                print(f"Saved CFG DOT file as: {out_path}")

            if cnsl.prompt_confirmation("Display CFG diagram?"):
                render_diagram(dot_str)
                print("Rendered CFG diagram.")

        elif mode == 3:
            # Reset global CFG builder state so each analysis run is clean
            cfg_build.FUNCTIONS.clear()
            cfg_build.CURRENT_FUNCTION = None

            cfg = make_cfg(ast)
            analysis_name = "Live Variables + Reached Uses"
            print(f"\nRunning {analysis_name} …")
            try:
                _analyses, annotations, ru_edges = run_all_analyses(cfg)
            except Exception as exc:
                print(f"Analysis failed: {exc}")
                continue

            print(
                f"Done. {len(annotations)} LV annotation node(s), "
                f"{len(ru_edges)} RU definition node(s)."
            )

            dot_str = analysis_to_dot(cfg, annotations, analysis_name, ru_edges)

            if cnsl.prompt_confirmation("\nPrint analysis reports to console?"):
                print_analysis_reports(cfg, _analyses, ru_edges)

            if cnsl.prompt_confirmation("\nExport annotated CFG as .dot file?"):
                default = f"{path.stem}_analysis.dot"
                filename = input(f"Filename [{default}]: ").strip()
                if not filename:
                    filename = default

                out_dir = Path(__file__).parent / "cfadots"
                out_dir.mkdir(exist_ok=True)
                out_path = out_dir / filename
                try:
                    with open(out_path, "w") as f:
                        f.write(dot_str)
                    print(f"Saved annotated DOT file as: {out_path}")
                except Exception as exc:
                    print(f"Could not save DOT file: {exc}")

            if cnsl.prompt_confirmation("Display annotated CFG diagram?"):
                render_diagram(dot_str)
                print("Rendered annotated CFG diagram.")
