import customtkinter as ctk
import tkinter as tk
from tkinter import filedialog, messagebox
import subprocess
import re
import threading
from pathlib import Path
import os
import tempfile

class TranslatorApp(ctk.CTk):
    def __init__(self):
        super().__init__()

        # Window configuration
        self.title("Ollama Translator - Русский переводчик")
        self.geometry("1000x700")

        # Set theme
        ctk.set_appearance_mode("dark")
        ctk.set_default_color_theme("blue")

        # Variables
        self.is_translating = False
        self.current_file = None

        # Create UI
        self.create_widgets()

        # Enable drag and drop (using tkinterdnd2 if available)
        try:
            self.setup_drag_drop()
        except:
            pass  # Drag and drop not available, skip it

    def create_widgets(self):
        # Main container
        self.grid_columnconfigure(0, weight=1)
        self.grid_rowconfigure(1, weight=1)

        # Header
        header_frame = ctk.CTkFrame(self, fg_color="transparent")
        header_frame.grid(row=0, column=0, padx=20, pady=(20, 10), sticky="ew")

        title_label = ctk.CTkLabel(
            header_frame,
            text="🌐 Ollama Translator",
            font=ctk.CTkFont(size=24, weight="bold")
        )
        title_label.pack(side="left")

        subtitle_label = ctk.CTkLabel(
            header_frame,
            text="Uncensored перевод на русский язык",
            font=ctk.CTkFont(size=12),
            text_color="gray"
        )
        subtitle_label.pack(side="left", padx=10)

        # Theme switch
        self.theme_switch = ctk.CTkSwitch(
            header_frame,
            text="🌙 Темная тема",
            command=self.toggle_theme,
            onvalue="dark",
            offvalue="light"
        )
        self.theme_switch.pack(side="right")
        self.theme_switch.select()

        # Main content area
        content_frame = ctk.CTkFrame(self)
        content_frame.grid(row=1, column=0, padx=20, pady=10, sticky="nsew")
        content_frame.grid_columnconfigure(0, weight=1)
        content_frame.grid_columnconfigure(1, weight=1)
        content_frame.grid_rowconfigure(1, weight=1)

        # Source text area
        source_label = ctk.CTkLabel(
            content_frame,
            text="📝 Исходный текст",
            font=ctk.CTkFont(size=14, weight="bold")
        )
        source_label.grid(row=0, column=0, padx=10, pady=(10, 5), sticky="w")

        self.source_text = ctk.CTkTextbox(
            content_frame,
            font=ctk.CTkFont(size=13),
            wrap="word"
        )
        self.source_text.grid(row=1, column=0, padx=(10, 5), pady=(0, 10), sticky="nsew")

        # Character counter for source
        self.source_counter = ctk.CTkLabel(
            content_frame,
            text="Символов: 0",
            font=ctk.CTkFont(size=10),
            text_color="gray"
        )
        self.source_counter.grid(row=2, column=0, padx=10, pady=(0, 5), sticky="w")

        # Bind text change event
        self.source_text.bind("<KeyRelease>", self.update_char_count)

        # Translation text area
        translation_label = ctk.CTkLabel(
            content_frame,
            text="🇷🇺 Перевод на русский",
            font=ctk.CTkFont(size=14, weight="bold")
        )
        translation_label.grid(row=0, column=1, padx=10, pady=(10, 5), sticky="w")

        self.translation_text = ctk.CTkTextbox(
            content_frame,
            font=ctk.CTkFont(size=13),
            wrap="word"
        )
        self.translation_text.grid(row=1, column=1, padx=(5, 10), pady=(0, 10), sticky="nsew")

        # Character counter for translation
        self.translation_counter = ctk.CTkLabel(
            content_frame,
            text="Символов: 0",
            font=ctk.CTkFont(size=10),
            text_color="gray"
        )
        self.translation_counter.grid(row=2, column=1, padx=10, pady=(0, 5), sticky="w")

        # Control buttons
        button_frame = ctk.CTkFrame(self, fg_color="transparent")
        button_frame.grid(row=2, column=0, padx=20, pady=10, sticky="ew")

        # Left side buttons
        left_buttons = ctk.CTkFrame(button_frame, fg_color="transparent")
        left_buttons.pack(side="left", fill="x", expand=True)

        self.open_btn = ctk.CTkButton(
            left_buttons,
            text="📁 Открыть файл",
            command=self.open_file,
            width=140,
            height=35
        )
        self.open_btn.pack(side="left", padx=5)

        self.paste_btn = ctk.CTkButton(
            left_buttons,
            text="📋 Вставить",
            command=self.paste_text,
            width=120,
            height=35
        )
        self.paste_btn.pack(side="left", padx=5)

        self.clear_btn = ctk.CTkButton(
            left_buttons,
            text="🗑️ Очистить",
            command=self.clear_all,
            width=120,
            height=35,
            fg_color="gray40",
            hover_color="gray30"
        )
        self.clear_btn.pack(side="left", padx=5)

        # Center button (Translate)
        self.translate_btn = ctk.CTkButton(
            button_frame,
            text="🔄 Перевести",
            command=self.translate_text,
            width=180,
            height=40,
            font=ctk.CTkFont(size=14, weight="bold"),
            fg_color="#1f6aa5",
            hover_color="#144870"
        )
        self.translate_btn.pack(side="left", padx=20)

        # Right side buttons
        right_buttons = ctk.CTkFrame(button_frame, fg_color="transparent")
        right_buttons.pack(side="right", fill="x", expand=True)

        self.copy_btn = ctk.CTkButton(
            right_buttons,
            text="📄 Копировать",
            command=self.copy_translation,
            width=140,
            height=35
        )
        self.copy_btn.pack(side="right", padx=5)

        self.save_btn = ctk.CTkButton(
            right_buttons,
            text="💾 Сохранить",
            command=self.save_translation,
            width=140,
            height=35
        )
        self.save_btn.pack(side="right", padx=5)

        # Status bar
        self.status_frame = ctk.CTkFrame(self, height=30)
        self.status_frame.grid(row=3, column=0, padx=20, pady=(0, 20), sticky="ew")

        self.status_label = ctk.CTkLabel(
            self.status_frame,
            text="✅ Готов к работе",
            font=ctk.CTkFont(size=11)
        )
        self.status_label.pack(side="left", padx=10, pady=5)

        self.progress_bar = ctk.CTkProgressBar(self.status_frame, width=200)
        self.progress_bar.pack(side="right", padx=10, pady=5)
        self.progress_bar.set(0)

        # Keyboard shortcuts
        self.bind("<Control-o>", lambda e: self.open_file())
        self.bind("<Control-v>", lambda e: self.paste_text())
        self.bind("<Control-s>", lambda e: self.save_translation())
        self.bind("<Control-Return>", lambda e: self.translate_text())
        self.bind("<F5>", lambda e: self.translate_text())

    def setup_drag_drop(self):
        """Setup drag and drop for files (requires tkinterdnd2)"""
        try:
            from tkinterdnd2 import DND_FILES, TkinterDnD
            # This would require tkinterdnd2 package
            # For now, we'll skip this feature
            pass
        except ImportError:
            # tkinterdnd2 not available, drag-drop disabled
            pass

    def on_drop(self, event):
        """Handle file drop"""
        files = self.tk.splitlist(event.data)
        if files:
            self.load_file(files[0])

    def toggle_theme(self):
        """Toggle between dark and light theme"""
        if self.theme_switch.get() == "dark":
            ctk.set_appearance_mode("dark")
        else:
            ctk.set_appearance_mode("light")

    def update_char_count(self, event=None):
        """Update character counter"""
        text = self.source_text.get("1.0", "end-1c")
        char_count = len(text)
        word_count = len(text.split())
        self.source_counter.configure(text=f"Символов: {char_count} | Слов: {word_count}")

    def update_translation_count(self):
        """Update translation character counter"""
        text = self.translation_text.get("1.0", "end-1c")
        char_count = len(text)
        word_count = len(text.split())
        self.translation_counter.configure(text=f"Символов: {char_count} | Слов: {word_count}")

    def open_file(self):
        """Open file dialog and load file"""
        file_path = filedialog.askopenfilename(
            title="Выберите файл для перевода",
            filetypes=[
                ("Текстовые файлы", "*.txt"),
                ("Markdown файлы", "*.md"),
                ("Все файлы", "*.*")
            ]
        )
        if file_path:
            self.load_file(file_path)

    def load_file(self, file_path):
        """Load file content"""
        try:
            # Try different encodings
            encodings = ['utf-8', 'cp1251', 'latin-1']
            content = None

            for encoding in encodings:
                try:
                    with open(file_path, 'r', encoding=encoding) as f:
                        content = f.read()
                    break
                except UnicodeDecodeError:
                    continue

            if content is None:
                messagebox.showerror("Ошибка", "Не удалось прочитать файл")
                return

            self.source_text.delete("1.0", "end")
            self.source_text.insert("1.0", content)
            self.current_file = file_path
            self.update_char_count()
            self.update_status(f"📁 Загружен: {Path(file_path).name}")

        except Exception as e:
            messagebox.showerror("Ошибка", f"Не удалось открыть файл:\n{str(e)}")

    def paste_text(self):
        """Paste text from clipboard"""
        try:
            clipboard_text = self.clipboard_get()
            self.source_text.delete("1.0", "end")
            self.source_text.insert("1.0", clipboard_text)
            self.update_char_count()
            self.update_status("📋 Текст вставлен из буфера обмена")
        except:
            messagebox.showwarning("Предупреждение", "Буфер обмена пуст")

    def clear_all(self):
        """Clear all text fields"""
        self.source_text.delete("1.0", "end")
        self.translation_text.delete("1.0", "end")
        self.current_file = None
        self.update_char_count()
        self.update_translation_count()
        self.update_status("🗑️ Поля очищены")

    def translate_text(self):
        """Translate text using ollama model"""
        if self.is_translating:
            return

        text = self.source_text.get("1.0", "end-1c").strip()
        if not text:
            messagebox.showwarning("Предупреждение", "Введите текст для перевода")
            return

        # Run translation in separate thread
        thread = threading.Thread(target=self._translate_worker, args=(text,))
        thread.daemon = True
        thread.start()

    def _translate_worker(self, text):
        """Worker thread for translation"""
        self.is_translating = True
        self.update_status("⏳ Перевод...")
        self.translate_btn.configure(state="disabled", text="⏳ Перевожу...")

        try:
            # Split text into chunks if it's too long
            chunk_size = 800  # Characters per chunk
            chunks = self._split_text_into_chunks(text, chunk_size)
            total_chunks = len(chunks)

            translated_parts = []

            for i, chunk in enumerate(chunks):
                # Update progress
                progress = (i + 1) / total_chunks
                self.progress_bar.set(progress)

                if total_chunks > 1:
                    self.after(0, lambda p=i+1, t=total_chunks:
                              self.update_status(f"⏳ Перевод части {p}/{t}..."))

                # Translate chunk
                result = subprocess.run(
                    ['ollama', 'run', 'translator', chunk],
                    capture_output=True,
                    text=True,
                    encoding='utf-8',
                    timeout=120
                )

                # Filter output
                output = result.stdout
                output = re.sub(r'<think>.*?</think>', '', output, flags=re.DOTALL)
                output = '\n'.join(line.strip() for line in output.split('\n') if line.strip())

                translated_parts.append(output)

            # Combine all parts
            final_translation = '\n\n'.join(translated_parts)

            # Update UI in main thread
            self.after(0, self._update_translation, final_translation)

        except subprocess.TimeoutExpired:
            self.after(0, lambda: messagebox.showerror("Ошибка", "Превышено время ожидания перевода"))
            self.after(0, lambda: self.update_status("❌ Ошибка: таймаут"))
        except Exception as e:
            self.after(0, lambda: messagebox.showerror("Ошибка", f"Ошибка перевода:\n{str(e)}"))
            self.after(0, lambda: self.update_status("❌ Ошибка перевода"))
        finally:
            self.is_translating = False
            self.progress_bar.set(0)
            self.translate_btn.configure(state="normal", text="🔄 Перевести")

    def _split_text_into_chunks(self, text, chunk_size=800):
        """Split text into chunks, trying to break at sentence boundaries"""
        if len(text) <= chunk_size:
            return [text]

        chunks = []
        current_chunk = ""

        # Split by sentences (simple approach)
        sentences = re.split(r'([.!?]\s+|\n\n)', text)

        for i in range(0, len(sentences), 2):
            sentence = sentences[i]
            separator = sentences[i + 1] if i + 1 < len(sentences) else ""

            # If adding this sentence exceeds chunk size, start new chunk
            if len(current_chunk) + len(sentence) + len(separator) > chunk_size and current_chunk:
                chunks.append(current_chunk.strip())
                current_chunk = sentence + separator
            else:
                current_chunk += sentence + separator

        # Add remaining text
        if current_chunk.strip():
            chunks.append(current_chunk.strip())

        # If no sentence boundaries found, split by character count
        if not chunks:
            chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]

        return chunks

    def _update_translation(self, translation):
        """Update translation text in UI"""
        self.translation_text.delete("1.0", "end")
        self.translation_text.insert("1.0", translation)
        self.update_translation_count()
        self.progress_bar.set(1.0)
        self.update_status("✅ Перевод завершен")

    def copy_translation(self):
        """Copy translation to clipboard"""
        translation = self.translation_text.get("1.0", "end-1c")
        if translation.strip():
            self.clipboard_clear()
            self.clipboard_append(translation)
            self.update_status("📄 Перевод скопирован в буфер обмена")
        else:
            messagebox.showwarning("Предупреждение", "Нет текста для копирования")

    def save_translation(self):
        """Save translation to file"""
        translation = self.translation_text.get("1.0", "end-1c")
        if not translation.strip():
            messagebox.showwarning("Предупреждение", "Нет текста для сохранения")
            return

        # Suggest filename
        if self.current_file:
            default_name = Path(self.current_file).stem + "_ru.txt"
        else:
            default_name = "translation_ru.txt"

        file_path = filedialog.asksaveasfilename(
            title="Сохранить перевод",
            defaultextension=".txt",
            initialfile=default_name,
            filetypes=[
                ("Текстовые файлы", "*.txt"),
                ("Markdown файлы", "*.md"),
                ("Все файлы", "*.*")
            ]
        )

        if file_path:
            try:
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write(translation)
                self.update_status(f"💾 Сохранено: {Path(file_path).name}")
                messagebox.showinfo("Успех", f"Перевод сохранен:\n{file_path}")
            except Exception as e:
                messagebox.showerror("Ошибка", f"Не удалось сохранить файл:\n{str(e)}")

    def update_status(self, message):
        """Update status bar message"""
        self.status_label.configure(text=message)


if __name__ == "__main__":
    app = TranslatorApp()
    app.mainloop()