# gui.py
import tomllib
import tomli_w
import imgui
import pygame
from imgui.integrations.pygame import PygameRenderer
from pathlib import Path
from tkinter import filedialog, Tk
import threading
import aggregate
import gemini

CONFIG_PATH = Path("config.toml")

def load_config() -> dict:
    with open(CONFIG_PATH, "rb") as f:
        return tomllib.load(f)

def save_config(config: dict):
    with open(CONFIG_PATH, "wb") as f:
        tomli_w.dump(config, f)

def hide_tk_root() -> Tk:
    root = Tk()
    root.withdraw()
    root.wm_attributes("-topmost", True)
    return root

class State:
    def __init__(self):
        self.config = load_config()

        self.namespace_buf = self.config["output"]["namespace"]
        self.output_dir_buf = self.config["output"]["output_dir"]
        self.files_base_dir_buf = self.config["files"]["base_dir"]
        self.screenshots_base_dir_buf = self.config.get("screenshots", {}).get("base_dir", ".")

        self.files = list(self.config["files"].get("paths", []))
        self.screenshots = list(self.config.get("screenshots", {}).get("paths", []))
        self.history = list(self.config.get("discussion", {}).get("history", []))

        self.history_edit_bufs = [h for h in self.history]

        self.gemini_input_buf = ""
        self.gemini_response = ""
        self.gemini_status = "idle"
        self.last_md = ""
        self.last_md_path = None
        self.send_thread = None

    def flush_to_config(self):
        self.config["output"]["namespace"] = self.namespace_buf
        self.config["output"]["output_dir"] = self.output_dir_buf
        self.config["files"]["base_dir"] = self.files_base_dir_buf
        if "screenshots" not in self.config:
            self.config["screenshots"] = {}
        self.config["screenshots"]["base_dir"] = self.screenshots_base_dir_buf
        self.config["files"]["paths"] = self.files
        self.config["screenshots"]["paths"] = self.screenshots
        self.config["discussion"] = {"history": self.history_edit_bufs}

    def sync_history_bufs(self):
        while len(self.history_edit_bufs) < len(self.history):
            self.history_edit_bufs.append("")
        self.history_edit_bufs = self.history_edit_bufs[:len(self.history)]

def draw_config_panel(state: State):
    imgui.begin("Config")

    changed, val = imgui.input_text("Namespace", state.namespace_buf, 256)
    if changed:
        state.namespace_buf = val

    changed, val = imgui.input_text("Output Dir", state.output_dir_buf, 512)
    if changed:
        state.output_dir_buf = val

    if imgui.button("Save Config"):
        state.flush_to_config()
        save_config(state.config)

    imgui.end()

def draw_files_panel(state: State):
    imgui.begin("Files")

    changed, val = imgui.input_text("Base Dir##files", state.files_base_dir_buf, 512)
    if changed:
        state.files_base_dir_buf = val

    if imgui.button("Browse Base Dir##files"):
        root = hide_tk_root()
        d = filedialog.askdirectory(title="Select Files Base Dir")
        root.destroy()
        if d:
            state.files_base_dir_buf = d

    imgui.separator()

    to_remove = None
    for i, f in enumerate(state.files):
        imgui.text(f)
        imgui.same_line()
        if imgui.button(f"Remove##file{i}"):
            to_remove = i
    if to_remove is not None:
        state.files.pop(to_remove)

    if imgui.button("Add File"):
        root = hide_tk_root()
        paths = filedialog.askopenfilenames(title="Select Files")
        root.destroy()
        for p in paths:
            if p not in state.files:
                state.files.append(p)

    if imgui.button("Add Wildcard##files"):
        root = hide_tk_root()
        d = filedialog.askdirectory(title="Select Dir for Wildcard")
        root.destroy()
        if d:
            state.files.append(str(Path(d) / "**" / "*"))

    imgui.end()

def draw_screenshots_panel(state: State):
    imgui.begin("Screenshots")

    changed, val = imgui.input_text("Base Dir##shots", state.screenshots_base_dir_buf, 512)
    if changed:
        state.screenshots_base_dir_buf = val

    if imgui.button("Browse Base Dir##shots"):
        root = hide_tk_root()
        d = filedialog.askdirectory(title="Select Screenshots Base Dir")
        root.destroy()
        if d:
            state.screenshots_base_dir_buf = d

    imgui.separator()

    to_remove = None
    for i, s in enumerate(state.screenshots):
        imgui.text(s)
        imgui.same_line()
        if imgui.button(f"Remove##shot{i}"):
            to_remove = i
    if to_remove is not None:
        state.screenshots.pop(to_remove)

    if imgui.button("Add Screenshot"):
        root = hide_tk_root()
        paths = filedialog.askopenfilenames(
            title="Select Screenshots",
            filetypes=[("Images", "*.png *.jpg *.jpeg *.gif *.bmp *.webp"), ("All", "*.*")]
        )
        root.destroy()
        for p in paths:
            if p not in state.screenshots:
                state.screenshots.append(p)

    imgui.end()

def draw_discussion_panel(state: State):
    imgui.begin("Discussion History")

    if imgui.button("Add Excerpt"):
        state.history.append("")
        state.history_edit_bufs.append("")

    imgui.separator()

    state.sync_history_bufs()
    to_remove = None
    for i in range(len(state.history)):
        imgui.text(f"Excerpt {i + 1}")
        imgui.same_line()
        if imgui.button(f"Remove##hist{i}"):
            to_remove = i
            continue
        changed, val = imgui.input_text_multiline(
            f"##histbuf{i}",
            state.history_edit_bufs[i],
            2048,
            width=-1,
            height=120
        )
        if changed:
            state.history_edit_bufs[i] = val
        imgui.separator()

    if to_remove is not None:
        state.history.pop(to_remove)
        state.history_edit_bufs.pop(to_remove)

    imgui.end()

def draw_gemini_panel(state: State):
    imgui.begin("Gemini")

    imgui.text(f"Status: {state.gemini_status}")

    if state.last_md_path:
        imgui.text(f"Last MD: {state.last_md_path}")

    imgui.separator()

    changed, val = imgui.input_text_multiline(
        "##gemini_input",
        state.gemini_input_buf,
        4096,
        width=-1,
        height=100
    )
    if changed:
        state.gemini_input_buf = val

    is_sending = state.send_thread is not None and state.send_thread.is_alive()

    if is_sending:
        imgui.begin_disabled()

    if imgui.button("Generate MD + Send"):
        state.flush_to_config()
        save_config(state.config)

        md, path = aggregate.run(state.config)
        state.last_md = md
        state.last_md_path = path
        state.gemini_status = "sending..."

        user_msg = state.gemini_input_buf

        def do_send():
            try:
                response = gemini.send(state.last_md, user_msg)
                state.gemini_response = response
                state.gemini_status = "done"
            except Exception as e:
                state.gemini_response = f"ERROR: {e}"
                state.gemini_status = "error"

        state.send_thread = threading.Thread(target=do_send, daemon=True)
        state.send_thread.start()

    if is_sending:
        imgui.end_disabled()

    imgui.same_line()

    if imgui.button("Generate MD Only"):
        state.flush_to_config()
        save_config(state.config)
        md, path = aggregate.run(state.config)
        state.last_md = md
        state.last_md_path = path
        state.gemini_status = "md generated"

    imgui.same_line()

    if imgui.button("Reset Session"):
        gemini.reset_session()
        state.gemini_status = "session reset"
        state.gemini_response = ""

    imgui.separator()
    imgui.text("Response:")
    imgui.input_text_multiline(
        "##gemini_response",
        state.gemini_response,
        max_value_length=65536,
        width=-1,
        height=-1,
        flags=imgui.INPUT_TEXT_READ_ONLY
    )

    imgui.end()

def main():
    pygame.init()
    screen = pygame.display.set_mode(
        (1280, 720),
        pygame.RESIZABLE
    )
    pygame.display.set_caption("manual slop")

    imgui.create_context()
    io = imgui.get_io()
    io.display_size = (1280, 720)

    renderer = PygameRenderer()

    state = State()
    clock = pygame.time.Clock()

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return
            if event.type == pygame.VIDEORESIZE:
                io.display_size = (event.w, event.h)
            renderer.process_event(event)

        imgui.new_frame()

        vp = imgui.get_main_viewport()
        imgui.set_next_window_position(vp.pos.x, vp.pos.y)
        imgui.set_next_window_size(vp.size.x, vp.size.y)
        imgui.begin(
            "##root",
            flags=(
                imgui.WINDOW_NO_TITLE_BAR
                | imgui.WINDOW_NO_RESIZE
                | imgui.WINDOW_NO_MOVE
                | imgui.WINDOW_NO_SCROLLBAR
                | imgui.WINDOW_NO_SAVED_SETTINGS
            )
        )

        avail_w, avail_h = imgui.get_content_region_available()
        col_w = avail_w / 4

        imgui.begin_child("##left", width=col_w, height=avail_h, border=True)
        draw_config_panel_inline(state)
        imgui.separator()
        draw_files_panel_inline(state)
        imgui.end_child()

        imgui.same_line()

        imgui.begin_child("##mid", width=col_w, height=avail_h, border=True)
        draw_screenshots_panel_inline(state)
        imgui.end_child()

        imgui.same_line()

        imgui.begin_child("##right_top", width=col_w, height=avail_h, border=True)
        draw_discussion_panel_inline(state)
        imgui.end_child()

        imgui.same_line()

        imgui.begin_child("##gemini", width=col_w, height=avail_h, border=True)
        draw_gemini_panel_inline(state)
        imgui.end_child()

        imgui.end()

        screen.fill((30, 30, 30))
        imgui.render()
        renderer.render(imgui.get_draw_data())
        pygame.display.flip()
        clock.tick(30)


def draw_config_panel_inline(state: State):
    imgui.text("Config")
    imgui.separator()

    changed, val = imgui.input_text("Namespace", state.namespace_buf, 256)
    if changed:
        state.namespace_buf = val

    changed, val = imgui.input_text("Output Dir", state.output_dir_buf, 512)
    if changed:
        state.output_dir_buf = val

    if imgui.button("Browse Output Dir"):
        root = hide_tk_root()
        d = filedialog.askdirectory(title="Select Output Dir")
        root.destroy()
        if d:
            state.output_dir_buf = d

    if imgui.button("Save Config"):
        state.flush_to_config()
        save_config(state.config)


def draw_files_panel_inline(state: State):
    imgui.text("Files")
    imgui.separator()

    changed, val = imgui.input_text("Base Dir##files", state.files_base_dir_buf, 512)
    if changed:
        state.files_base_dir_buf = val

    if imgui.button("Browse Base Dir##files"):
        root = hide_tk_root()
        d = filedialog.askdirectory(title="Select Files Base Dir")
        root.destroy()
        if d:
            state.files_base_dir_buf = d

    imgui.separator()

    to_remove = None
    for i, f in enumerate(state.files):
        imgui.text(f[:40] + "..." if len(f) > 40 else f)
        imgui.same_line()
        if imgui.button(f"x##file{i}"):
            to_remove = i
    if to_remove is not None:
        state.files.pop(to_remove)

    if imgui.button("Add File(s)##files"):
        root = hide_tk_root()
        paths = filedialog.askopenfilenames(title="Select Files")
        root.destroy()
        for p in paths:
            if p not in state.files:
                state.files.append(p)

    if imgui.button("Add Wildcard##files"):
        root = hide_tk_root()
        d = filedialog.askdirectory(title="Select Dir for Wildcard")
        root.destroy()
        if d:
            state.files.append(str(Path(d) / "**" / "*"))


def draw_screenshots_panel_inline(state: State):
    imgui.text("Screenshots")
    imgui.separator()

    changed, val = imgui.input_text("Base Dir##shots", state.screenshots_base_dir_buf, 512)
    if changed:
        state.screenshots_base_dir_buf = val

    if imgui.button("Browse Base Dir##shots"):
        root = hide_tk_root()
        d = filedialog.askdirectory(title="Select Screenshots Base Dir")
        root.destroy()
        if d:
            state.screenshots_base_dir_buf = d

    imgui.separator()

    to_remove = None
    for i, s in enumerate(state.screenshots):
        imgui.text(s[:40] + "..." if len(s) > 40 else s)
        imgui.same_line()
        if imgui.button(f"x##shot{i}"):
            to_remove = i
    if to_remove is not None:
        state.screenshots.pop(to_remove)

    if imgui.button("Add Screenshot(s)"):
        root = hide_tk_root()
        paths = filedialog.askopenfilenames(
            title="Select Screenshots",
            filetypes=[("Images", "*.png *.jpg *.jpeg *.gif *.bmp *.webp"), ("All", "*.*")]
        )
        root.destroy()
        for p in paths:
            if p not in state.screenshots:
                state.screenshots.append(p)


def draw_discussion_panel_inline(state: State):
    imgui.text("Discussion History")
    imgui.separator()

    if imgui.button("Add Excerpt"):
        state.history.append("")
        state.history_edit_bufs.append("")

    imgui.separator()

    state.sync_history_bufs()
    to_remove = None
    for i in range(len(state.history)):
        imgui.text(f"Excerpt {i + 1}")
        imgui.same_line()
        if imgui.button(f"x##hist{i}"):
            to_remove = i
            continue
        changed, val = imgui.input_text_multiline(
            f"##histbuf{i}",
            state.history_edit_bufs[i],
            2048,
            width=-1,
            height=120
        )
        if changed:
            state.history_edit_bufs[i] = val
        imgui.separator()

    if to_remove is not None:
        state.history.pop(to_remove)
        state.history_edit_bufs.pop(to_remove)


def draw_gemini_panel_inline(state: State):
    imgui.text("Gemini")
    imgui.separator()
    imgui.text(f"Status: {state.gemini_status}")

    if state.last_md_path:
        imgui.text(f"Last MD: {str(state.last_md_path)[:40]}")

    imgui.separator()

    changed, val = imgui.input_text_multiline(
        "##gemini_input",
        state.gemini_input_buf,
        4096,
        width=-1,
        height=120
    )
    if changed:
        state.gemini_input_buf = val

    is_sending = state.send_thread is not None and state.send_thread.is_alive()

    if is_sending:
        imgui.begin_disabled()

    if imgui.button("Generate + Send"):
        state.flush_to_config()
        save_config(state.config)
        md, path = aggregate.run(state.config)
        state.last_md = md
        state.last_md_path = path
        state.gemini_status = "sending..."
        user_msg = state.gemini_input_buf

        def do_send():
            try:
                response = gemini.send(state.last_md, user_msg)
                state.gemini_response = response
                state.gemini_status = "done"
            except Exception as e:
                state.gemini_response = f"ERROR: {e}"
                state.gemini_status = "error"

        state.send_thread = threading.Thread(target=do_send, daemon=True)
        state.send_thread.start()

    if is_sending:
        imgui.end_disabled()

    imgui.same_line()

    if imgui.button("MD Only"):
        state.flush_to_config()
        save_config(state.config)
        md, path = aggregate.run(state.config)
        state.last_md = md
        state.last_md_path = path
        state.gemini_status = "md generated"

    imgui.same_line()

    if imgui.button("Reset"):
        gemini.reset_session()
        state.gemini_status = "session reset"
        state.gemini_response = ""

    imgui.separator()
    imgui.text("Response:")

    _, _ = imgui.input_text_multiline(
        "##gemini_response",
        state.gemini_response,
        max_value_length=65536,
        width=-1,
        height=-1,
        flags=imgui.INPUT_TEXT_READ_ONLY
    )


if __name__ == "__main__":
    main()