import ctypes
import json
import os
import threading
import time
from dataclasses import dataclass, field
from typing import List, Optional

import keyboard
import pyperclip
import requests
import tkinter as tk
from tkinter import ttk, messagebox

import pystray
from PIL import Image, ImageDraw

APP_DIR = os.path.join(os.path.expanduser("~"), ".ollama_tray_rewriter")
CONFIG_PATH = os.path.join(APP_DIR, "config.json")

DEFAULT_CONFIG = {
    "hotkey": "ctrl+shift+space",
    "ollama_url": "http://127.0.0.1:11434",
    "model": "llama3.1:8b",
    "prompts": [
        {"name": "Klarer formulieren", "text": "Formuliere den Text klarer und prägnanter, Sinn behalten."},
        {"name": "Professionell", "text": "Überarbeite den Text in professionellem, freundlichem Ton."}
    ]
}


user32 = ctypes.windll.user32


def get_foreground_window() -> int:
    return int(user32.GetForegroundWindow())


def focus_window(hwnd: int):
    if hwnd:
        user32.SetForegroundWindow(hwnd)


def ensure_app_dir():
    os.makedirs(APP_DIR, exist_ok=True)


def load_config() -> dict:
    ensure_app_dir()
    if not os.path.exists(CONFIG_PATH):
        save_config(DEFAULT_CONFIG)
        return dict(DEFAULT_CONFIG)
    try:
        with open(CONFIG_PATH, "r", encoding="utf-8") as f:
            data = json.load(f)
        for k, v in DEFAULT_CONFIG.items():
            if k not in data:
                data[k] = v
        return data
    except Exception:
        return dict(DEFAULT_CONFIG)


def save_config(cfg: dict):
    ensure_app_dir()
    with open(CONFIG_PATH, "w", encoding="utf-8") as f:
        json.dump(cfg, f, ensure_ascii=False, indent=2)


def call_ollama(base_url: str, model: str, prompt_instruction: str, source: str) -> str:
    url = base_url.rstrip("/") + "/api/generate"
    full_prompt = f"{prompt_instruction}\n\nTEXT:\n{source}"
    payload = {
        "model": model,
        "prompt": full_prompt,
        "stream": False
    }
    r = requests.post(url, json=payload, timeout=120)
    r.raise_for_status()
    data = r.json()
    return data.get("response", "")


class RewriteWindow:
    def __init__(self, cfg: dict):
        self.cfg = cfg
        self.previous_hwnd = 0

        self.root = tk.Tk()
        self.root.title("Ollama Tray Rewriter")
        self.root.geometry("980x720")
        self.root.withdraw()

        top = ttk.Frame(self.root, padding=8)
        top.pack(fill="x")

        ttk.Label(top, text="Ollama URL:").grid(row=0, column=0, sticky="w")
        self.url_var = tk.StringVar(value=self.cfg.get("ollama_url", ""))
        ttk.Entry(top, textvariable=self.url_var, width=36).grid(row=0, column=1, padx=6)

        ttk.Label(top, text="Modell:").grid(row=0, column=2, sticky="w")
        self.model_var = tk.StringVar(value=self.cfg.get("model", ""))
        ttk.Entry(top, textvariable=self.model_var, width=24).grid(row=0, column=3, padx=6)

        ttk.Label(top, text="Hotkey:").grid(row=0, column=4, sticky="w")
        self.hotkey_var = tk.StringVar(value=self.cfg.get("hotkey", "ctrl+shift+space"))
        ttk.Entry(top, textvariable=self.hotkey_var, width=18).grid(row=0, column=5, padx=6)

        ttk.Button(top, text="Speichern", command=self.save_settings).grid(row=0, column=6, padx=6)

        prompt_bar = ttk.Frame(self.root, padding=(8, 0, 8, 8))
        prompt_bar.pack(fill="x")

        self.prompt_names = [p.get("name", "Prompt") for p in self.cfg.get("prompts", [])]
        self.prompt_var = tk.StringVar(value=self.prompt_names[0] if self.prompt_names else "")
        ttk.Label(prompt_bar, text="Prompt:").pack(side="left")
        self.prompt_combo = ttk.Combobox(prompt_bar, values=self.prompt_names, textvariable=self.prompt_var, state="readonly", width=34)
        self.prompt_combo.pack(side="left", padx=6)
        self.prompt_combo.bind("<<ComboboxSelected>>", lambda e: self.load_prompt_text())

        ttk.Button(prompt_bar, text="Neu", command=self.add_prompt).pack(side="left", padx=4)
        ttk.Button(prompt_bar, text="Löschen", command=self.delete_prompt).pack(side="left", padx=4)

        ttk.Label(prompt_bar, text="Prompt-Text:").pack(side="left", padx=(20, 4))

        self.prompt_text = tk.Text(self.root, height=5, wrap="word")
        self.prompt_text.pack(fill="x", padx=8)

        panes = ttk.Panedwindow(self.root, orient="vertical")
        panes.pack(fill="both", expand=True, padx=8, pady=8)

        src_frame = ttk.Labelframe(panes, text="Quelle")
        self.source_text = tk.Text(src_frame, wrap="word")
        self.source_text.pack(fill="both", expand=True)
        panes.add(src_frame, weight=1)

        dst_frame = ttk.Labelframe(panes, text="Ziel")
        self.target_text = tk.Text(dst_frame, wrap="word")
        self.target_text.pack(fill="both", expand=True)
        panes.add(dst_frame, weight=1)

        btns = ttk.Frame(self.root, padding=8)
        btns.pack(fill="x")

        self.status_var = tk.StringVar(value="Bereit")
        ttk.Label(btns, textvariable=self.status_var).pack(side="left")

        ttk.Button(btns, text="Bearbeiten", command=self.rewrite_now).pack(side="right", padx=4)
        ttk.Button(btns, text="Ins Clipboard", command=self.copy_and_close).pack(side="right", padx=4)
        ttk.Button(btns, text="Einfügen", command=self.paste_back_and_close).pack(side="right", padx=4)
        ttk.Button(btns, text="Schließen", command=self.hide).pack(side="right", padx=4)

        self.load_prompt_text()

    def save_settings(self):
        self.cfg["ollama_url"] = self.url_var.get().strip()
        self.cfg["model"] = self.model_var.get().strip()
        self.cfg["hotkey"] = self.hotkey_var.get().strip() or "ctrl+shift+space"

        prompts = self.cfg.get("prompts", [])
        idx = self.current_prompt_index()
        if idx is not None:
            prompts[idx]["text"] = self.prompt_text.get("1.0", "end").strip()
        self.cfg["prompts"] = prompts
        save_config(self.cfg)
        self.status_var.set("Gespeichert")

    def current_prompt_index(self) -> Optional[int]:
        name = self.prompt_var.get()
        for i, p in enumerate(self.cfg.get("prompts", [])):
            if p.get("name") == name:
                return i
        return None

    def load_prompt_text(self):
        idx = self.current_prompt_index()
        self.prompt_text.delete("1.0", "end")
        if idx is not None:
            self.prompt_text.insert("1.0", self.cfg["prompts"][idx].get("text", ""))

    def add_prompt(self):
        name = f"Prompt {len(self.cfg.get('prompts', [])) + 1}"
        self.cfg.setdefault("prompts", []).append({"name": name, "text": ""})
        self.refresh_prompt_combo(select=name)

    def delete_prompt(self):
        idx = self.current_prompt_index()
        if idx is None:
            return
        self.cfg["prompts"].pop(idx)
        self.refresh_prompt_combo()

    def refresh_prompt_combo(self, select: Optional[str] = None):
        self.prompt_names = [p.get("name", "Prompt") for p in self.cfg.get("prompts", [])]
        self.prompt_combo["values"] = self.prompt_names
        if self.prompt_names:
            self.prompt_var.set(select if select in self.prompt_names else self.prompt_names[0])
        else:
            self.prompt_var.set("")
        self.load_prompt_text()

    def show_with_source(self, source_text: str, previous_hwnd: int):
        self.previous_hwnd = previous_hwnd
        self.source_text.delete("1.0", "end")
        self.source_text.insert("1.0", source_text)
        self.target_text.delete("1.0", "end")

        self.root.deiconify()
        self.root.lift()
        self.root.attributes('-topmost', True)
        self.root.after(300, lambda: self.root.attributes('-topmost', False))
        self.status_var.set("Text übernommen")

    def hide(self):
        self.root.withdraw()

    def rewrite_now(self):
        idx = self.current_prompt_index()
        if idx is None:
            messagebox.showwarning("Prompt fehlt", "Bitte Prompt anlegen/auswählen.")
            return

        source = self.source_text.get("1.0", "end").strip()
        prompt_instruction = self.prompt_text.get("1.0", "end").strip()
        if not source or not prompt_instruction:
            messagebox.showwarning("Fehlt", "Quelle und Prompt dürfen nicht leer sein.")
            return

        self.status_var.set("Bearbeite via Ollama…")

        def worker():
            try:
                out = call_ollama(self.url_var.get().strip(), self.model_var.get().strip(), prompt_instruction, source)
                self.root.after(0, lambda: self._set_target(out))
            except Exception as e:
                self.root.after(0, lambda: messagebox.showerror("Ollama Fehler", str(e)))
                self.root.after(0, lambda: self.status_var.set("Fehler"))

        threading.Thread(target=worker, daemon=True).start()

    def _set_target(self, txt: str):
        self.target_text.delete("1.0", "end")
        self.target_text.insert("1.0", txt)
        self.status_var.set("Fertig")

    def copy_and_close(self):
        txt = self.target_text.get("1.0", "end").strip()
        pyperclip.copy(txt)
        self.hide()

    def paste_back_and_close(self):
        txt = self.target_text.get("1.0", "end").strip()
        pyperclip.copy(txt)
        self.hide()
        if self.previous_hwnd:
            time.sleep(0.12)
            focus_window(self.previous_hwnd)
            time.sleep(0.08)
            keyboard.send("ctrl+v")


class TrayApp:
    def __init__(self):
        self.cfg = load_config()
        self.window = RewriteWindow(self.cfg)
        self.icon = self._make_tray_icon()
        self.hotkey_handle = None
        self.register_hotkey(self.cfg.get("hotkey", "ctrl+shift+space"))

    def _make_image(self):
        img = Image.new("RGB", (64, 64), color=(30, 41, 59))
        d = ImageDraw.Draw(img)
        d.rectangle((10, 10, 54, 54), outline=(96, 165, 250), width=4)
        d.rectangle((22, 22, 42, 42), fill=(14, 165, 233))
        return img

    def _make_tray_icon(self):
        menu = pystray.Menu(
            pystray.MenuItem("Fenster öffnen", self._on_open),
            pystray.MenuItem("Beenden", self._on_quit),
        )
        return pystray.Icon("ollama_rewriter", self._make_image(), "Ollama Rewriter", menu)

    def register_hotkey(self, hotkey: str):
        if self.hotkey_handle is not None:
            try:
                keyboard.remove_hotkey(self.hotkey_handle)
            except Exception:
                pass
        self.hotkey_handle = keyboard.add_hotkey(hotkey, self.on_hotkey)

    def on_hotkey(self):
        prev = get_foreground_window()
        time.sleep(0.05)
        keyboard.send("ctrl+a")
        time.sleep(0.05)
        keyboard.send("ctrl+c")
        time.sleep(0.12)
        text = pyperclip.paste() or ""
        self.window.root.after(0, lambda: self.window.show_with_source(text, prev))

    def _on_open(self):
        self.window.root.after(0, lambda: self.window.show_with_source("", 0))

    def _on_quit(self):
        self.icon.stop()
        self.window.root.after(0, self.window.root.destroy)

    def run(self):
        t = threading.Thread(target=self.icon.run, daemon=True)
        t.start()
        self.window.root.mainloop()


if __name__ == "__main__":
    app = TrayApp()
    app.run()