import triplayacc as yacc
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 helpers.console as cnsl
import os

import matplotlib
matplotlib.use("TkAgg")

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

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

# Exports the AST as a DOT file
def export_dot_file(dot_string: str, filename: str):
    try:
        Path(filename).write_text(dot_string, encoding="utf-8")
        print(f"Saved DOT file as: {filename}")
    except Exception as e:
        print(f"Could not save DOT file: {e}")

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

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

        if mode == 2:
            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"
                cnsl = input(f"Filename [{default}]: ").strip()
                if not cnsl:
                    cnsl = default
                export_dot_file(dot_str, cnsl)

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

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

            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}")

            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()