# gui.py
import dearpygui.dearpygui as dpg
import tomllib
import tomli_w
import threading
from pathlib import Path
from tkinter import filedialog, Tk
import aggregate
import ai_client

CONFIG_PATH = Path("config.toml")
PROVIDERS = ["gemini", "anthropic"]

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 App:
    def __init__(self):
        self.config = load_config()
        self.files: list[str] = list(self.config["files"].get("paths", []))
        self.screenshots: list[str] = list(self.config.get("screenshots", {}).get("paths", []))
        self.history: list[str] = list(self.config.get("discussion", {}).get("history", []))

        ai_cfg = self.config.get("ai", {})
        self.current_provider: str = ai_cfg.get("provider", "gemini")
        self.current_model: str = ai_cfg.get("model", "gemini-2.0-flash")
        self.available_models: list[str] = []

        self.ai_status = "idle"
        self.ai_response = ""
        self.last_md = ""
        self.last_md_path: Path | None = None
        self.send_thread: threading.Thread | None = None
        self.models_thread: threading.Thread | None = None

        ai_client.set_provider(self.current_provider, self.current_model)

    # ------------------------------------------------------------------ helpers

    def _flush_to_config(self):
        self.config["output"]["namespace"] = dpg.get_value("namespace")
        self.config["output"]["output_dir"] = dpg.get_value("output_dir")
        self.config["files"]["base_dir"] = dpg.get_value("files_base_dir")
        if "screenshots" not in self.config:
            self.config["screenshots"] = {}
        self.config["screenshots"]["base_dir"] = dpg.get_value("shots_base_dir")
        self.config["files"]["paths"] = self.files
        self.config["screenshots"]["paths"] = self.screenshots
        raw = dpg.get_value("discussion_box")
        self.history = [s.strip() for s in raw.split("---") if s.strip()]
        self.config["discussion"] = {"history": self.history}
        self.config["ai"] = {
            "provider": self.current_provider,
            "model": self.current_model
        }

    def _do_generate(self) -> tuple[str, Path]:
        self._flush_to_config()
        save_config(self.config)
        return aggregate.run(self.config)

    def _update_status(self, status: str):
        self.ai_status = status
        if dpg.does_item_exist("ai_status"):
            dpg.set_value("ai_status", f"Status: {status}")

    def _update_response(self, text: str):
        self.ai_response = text
        if dpg.does_item_exist("ai_response"):
            dpg.set_value("ai_response", text)

    def _rebuild_files_list(self):
        if not dpg.does_item_exist("files_scroll"):
            return
        dpg.delete_item("files_scroll", children_only=True)
        for i, f in enumerate(self.files):
            with dpg.group(horizontal=True, parent="files_scroll"):
                dpg.add_button(
                    label="x",
                    width=24,
                    callback=self._make_remove_file_cb(i)
                )
                dpg.add_text(f)

    def _rebuild_shots_list(self):
        if not dpg.does_item_exist("shots_scroll"):
            return
        dpg.delete_item("shots_scroll", children_only=True)
        for i, s in enumerate(self.screenshots):
            with dpg.group(horizontal=True, parent="shots_scroll"):
                dpg.add_button(
                    label="x",
                    width=24,
                    callback=self._make_remove_shot_cb(i)
                )
                dpg.add_text(s)

    def _rebuild_models_list(self):
        if not dpg.does_item_exist("model_listbox"):
            return
        dpg.configure_item("model_listbox", items=self.available_models)
        if self.current_model in self.available_models:
            dpg.set_value("model_listbox", self.current_model)
        elif self.available_models:
            self.current_model = self.available_models[0]
            dpg.set_value("model_listbox", self.current_model)
            ai_client.set_provider(self.current_provider, self.current_model)

    def _make_remove_file_cb(self, idx: int):
        def cb():
            if idx < len(self.files):
                self.files.pop(idx)
                self._rebuild_files_list()
        return cb

    def _make_remove_shot_cb(self, idx: int):
        def cb():
            if idx < len(self.screenshots):
                self.screenshots.pop(idx)
                self._rebuild_shots_list()
        return cb

    def _fetch_models(self, provider: str):
        self._update_status("fetching models...")
        def do_fetch():
            try:
                models = ai_client.list_models(provider)
                self.available_models = models
                self._rebuild_models_list()
                self._update_status(f"models loaded: {len(models)}")
            except Exception as e:
                self._update_status(f"model fetch error: {e}")
        self.models_thread = threading.Thread(target=do_fetch, daemon=True)
        self.models_thread.start()

    # ---------------------------------------------------------------- callbacks

    def cb_browse_output(self):
        root = hide_tk_root()
        d = filedialog.askdirectory(title="Select Output Dir")
        root.destroy()
        if d:
            dpg.set_value("output_dir", d)

    def cb_save_config(self):
        self._flush_to_config()
        save_config(self.config)
        self._update_status("config saved")

    def cb_browse_files_base(self):
        root = hide_tk_root()
        d = filedialog.askdirectory(title="Select Files Base Dir")
        root.destroy()
        if d:
            dpg.set_value("files_base_dir", d)

    def cb_add_files(self):
        root = hide_tk_root()
        paths = filedialog.askopenfilenames(title="Select Files")
        root.destroy()
        for p in paths:
            if p not in self.files:
                self.files.append(p)
        self._rebuild_files_list()

    def cb_add_wildcard(self):
        root = hide_tk_root()
        d = filedialog.askdirectory(title="Select Dir for Wildcard")
        root.destroy()
        if d:
            self.files.append(str(Path(d) / "**" / "*"))
        self._rebuild_files_list()

    def cb_browse_shots_base(self):
        root = hide_tk_root()
        d = filedialog.askdirectory(title="Select Screenshots Base Dir")
        root.destroy()
        if d:
            dpg.set_value("shots_base_dir", d)

    def cb_add_shots(self):
        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 self.screenshots:
                self.screenshots.append(p)
        self._rebuild_shots_list()

    def cb_add_excerpt(self):
        current = dpg.get_value("discussion_box")
        dpg.set_value("discussion_box", current + "\n---\n")

    def cb_clear_discussion(self):
        dpg.set_value("discussion_box", "")

    def cb_save_discussion(self):
        self._flush_to_config()
        save_config(self.config)
        self._update_status("discussion saved")

    def cb_md_only(self):
        try:
            md, path = self._do_generate()
            self.last_md = md
            self.last_md_path = path
            self._update_status(f"md written: {path.name}")
        except Exception as e:
            self._update_status(f"error: {e}")

    def cb_reset_session(self):
        ai_client.reset_session()
        self._update_status("session reset")
        self._update_response("")

    def cb_generate_send(self):
        if self.send_thread and self.send_thread.is_alive():
            return
        try:
            md, path = self._do_generate()
            self.last_md = md
            self.last_md_path = path
        except Exception as e:
            self._update_status(f"generate error: {e}")
            return
        self._update_status("sending...")
        user_msg = dpg.get_value("ai_input")

        def do_send():
            try:
                response = ai_client.send(self.last_md, user_msg)
                self._update_response(response)
                self._update_status("done")
            except Exception as e:
                self._update_response(f"ERROR: {e}")
                self._update_status("error")

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

    def cb_provider_changed(self, sender, app_data):
        self.current_provider = app_data
        ai_client.reset_session()
        ai_client.set_provider(self.current_provider, self.current_model)
        self.available_models = []
        self._rebuild_models_list()
        self._fetch_models(self.current_provider)

    def cb_model_changed(self, sender, app_data):
        if app_data:
            self.current_model = app_data
            ai_client.reset_session()
            ai_client.set_provider(self.current_provider, self.current_model)
            self._update_status(f"model set: {self.current_model}")

    def cb_fetch_models(self):
        self._fetch_models(self.current_provider)

    # ---------------------------------------------------------------- build ui

    def _build_ui(self):

        with dpg.window(
            label="Config",
            tag="win_config",
            pos=(8, 8),
            width=400,
            height=200,
            no_close=True
        ):
            dpg.add_text("Namespace")
            dpg.add_input_text(
                tag="namespace",
                default_value=self.config["output"]["namespace"],
                width=-1
            )
            dpg.add_text("Output Dir")
            dpg.add_input_text(
                tag="output_dir",
                default_value=self.config["output"]["output_dir"],
                width=-1
            )
            with dpg.group(horizontal=True):
                dpg.add_button(label="Browse Output Dir", callback=self.cb_browse_output)
                dpg.add_button(label="Save Config", callback=self.cb_save_config)

        with dpg.window(
            label="Files",
            tag="win_files",
            pos=(8, 216),
            width=400,
            height=500,
            no_close=True
        ):
            dpg.add_text("Base Dir")
            with dpg.group(horizontal=True):
                dpg.add_input_text(
                    tag="files_base_dir",
                    default_value=self.config["files"]["base_dir"],
                    width=-220
                )
                dpg.add_button(label="Browse##filesbase", callback=self.cb_browse_files_base)
            dpg.add_separator()
            dpg.add_text("Paths")
            with dpg.child_window(tag="files_scroll", height=-64, border=True):
                pass
            self._rebuild_files_list()
            dpg.add_separator()
            with dpg.group(horizontal=True):
                dpg.add_button(label="Add File(s)", callback=self.cb_add_files)
                dpg.add_button(label="Add Wildcard", callback=self.cb_add_wildcard)

        with dpg.window(
            label="Screenshots",
            tag="win_screenshots",
            pos=(416, 8),
            width=400,
            height=500,
            no_close=True
        ):
            dpg.add_text("Base Dir")
            with dpg.group(horizontal=True):
                dpg.add_input_text(
                    tag="shots_base_dir",
                    default_value=self.config.get("screenshots", {}).get("base_dir", "."),
                    width=-220
                )
                dpg.add_button(label="Browse##shotsbase", callback=self.cb_browse_shots_base)
            dpg.add_separator()
            dpg.add_text("Paths")
            with dpg.child_window(tag="shots_scroll", height=-48, border=True):
                pass
            self._rebuild_shots_list()
            dpg.add_separator()
            dpg.add_button(label="Add Screenshot(s)", callback=self.cb_add_shots)

        with dpg.window(
            label="Discussion History",
            tag="win_discussion",
            pos=(824, 8),
            width=400,
            height=500,
            no_close=True
        ):
            dpg.add_input_text(
                tag="discussion_box",
                default_value="\n---\n".join(self.history),
                multiline=True,
                width=-1,
                height=-64
            )
            dpg.add_separator()
            with dpg.group(horizontal=True):
                dpg.add_button(label="Add Separator", callback=self.cb_add_excerpt)
                dpg.add_button(label="Clear", callback=self.cb_clear_discussion)
                dpg.add_button(label="Save", callback=self.cb_save_discussion)

        with dpg.window(
            label="Provider",
            tag="win_provider",
            pos=(1232, 8),
            width=420,
            height=280,
            no_close=True
        ):
            dpg.add_text("Provider")
            dpg.add_combo(
                tag="provider_combo",
                items=PROVIDERS,
                default_value=self.current_provider,
                width=-1,
                callback=self.cb_provider_changed
            )
            dpg.add_separator()
            with dpg.group(horizontal=True):
                dpg.add_text("Model")
                dpg.add_button(label="Fetch Models", callback=self.cb_fetch_models)
            dpg.add_listbox(
                tag="model_listbox",
                items=self.available_models,
                default_value=self.current_model,
                width=-1,
                num_items=6,
                callback=self.cb_model_changed
            )
            dpg.add_separator()
            dpg.add_text("Status: idle", tag="ai_status")

        with dpg.window(
            label="Message",
            tag="win_message",
            pos=(1232, 296),
            width=420,
            height=280,
            no_close=True
        ):
            dpg.add_input_text(
                tag="ai_input",
                multiline=True,
                width=-1,
                height=-64
            )
            dpg.add_separator()
            with dpg.group(horizontal=True):
                dpg.add_button(label="Gen + Send", callback=self.cb_generate_send)
                dpg.add_button(label="MD Only", callback=self.cb_md_only)
                dpg.add_button(label="Reset", callback=self.cb_reset_session)

        with dpg.window(
            label="Response",
            tag="win_response",
            pos=(1232, 584),
            width=420,
            height=400,
            no_close=True
        ):
            dpg.add_input_text(
                tag="ai_response",
                multiline=True,
                readonly=True,
                width=-1,
                height=-1
            )

    def run(self):
        dpg.create_context()
        dpg.configure_app(docking=True, docking_space=True)
        dpg.create_viewport(title="manual slop", width=1600, height=900)
        dpg.setup_dearpygui()
        dpg.show_viewport()
        dpg.maximize_viewport()
        self._build_ui()
        self._fetch_models(self.current_provider)

        while dpg.is_dearpygui_running():
            dpg.render_dearpygui_frame()

        dpg.destroy_context()


def main():
    app = App()
    app.run()


if __name__ == "__main__":
    main()