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 lib.console as cnsl
import syntax
import triplayacc as yacc
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}")


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", "Exit"])

        if mode == 3:
            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 / filename
                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.")