import triplayacc as yacc
import cfg_build
import syntax

from pathlib import Path
from graphviz import Source
from vistram.tram import *

import tkinter as tk
from vistram.vistram import MachineUI

import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import lib.console as cnsl
import os

from cfg.CFG import CFG
from cfg.CFG_Node import (CFG_START, CFG_END)

import matplotlib
matplotlib.use("TkAgg")

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

def make_cfg(ast):
    start = CFG_START()
    end   = CFG_END()

    last = ast.cfa(start, end)
    if last is not None:
        last.add_child(end)

    return CFG(start, end)

# 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.")




