#!/usr/bin/env python3
"""
MPQ Inspector - A GUI tool to inspect MPQ archives
Uses tkinter for the GUI and PyStorm for MPQ operations
"""

import tkinter as tk
from tkinter import ttk, filedialog, messagebox, scrolledtext
from pathlib import Path
import sys
from typing import Optional

try:
    from pystorm import MPQArchive, StormLibError
except ImportError:
    print("Error: PyStorm not installed. Please run: pip install -e .")
    sys.exit(1)


class MPQInspectorApp:
    """Main application window for MPQ Inspector"""
    
    def __init__(self, root):
        self.root = root
        self.root.title("MPQ Inspector - PyStorm Demo")
        self.root.geometry("900x700")
        
        self.current_archive: Optional[MPQArchive] = None
        self.current_path: Optional[str] = None
        
        self.setup_ui()
        
    def setup_ui(self):
        """Setup the user interface"""
        # Menu bar
        menubar = tk.Menu(self.root)
        self.root.config(menu=menubar)
        
        file_menu = tk.Menu(menubar, tearoff=0)
        menubar.add_cascade(label="File", menu=file_menu)
        file_menu.add_command(label="Open MPQ...", command=self.open_archive, accelerator="Ctrl+O")
        file_menu.add_command(label="Close Archive", command=self.close_archive, accelerator="Ctrl+W")
        file_menu.add_separator()
        file_menu.add_command(label="Exit", command=self.root.quit, accelerator="Ctrl+Q")
        
        help_menu = tk.Menu(menubar, tearoff=0)
        menubar.add_cascade(label="Help", menu=help_menu)
        help_menu.add_command(label="About", command=self.show_about)
        
        # Bind keyboard shortcuts
        self.root.bind('<Control-o>', lambda e: self.open_archive())
        self.root.bind('<Control-w>', lambda e: self.close_archive())
        self.root.bind('<Control-q>', lambda e: self.root.quit())
        
        # Main container
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(0, weight=1)
        main_frame.rowconfigure(2, weight=1)
        
        # Top bar - File selection
        top_frame = ttk.LabelFrame(main_frame, text="Archive", padding="5")
        top_frame.grid(row=0, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
        top_frame.columnconfigure(1, weight=1)
        
        ttk.Label(top_frame, text="File:").grid(row=0, column=0, sticky=tk.W, padx=(0, 5))
        
        self.file_path_var = tk.StringVar(value="No archive loaded")
        file_entry = ttk.Entry(top_frame, textvariable=self.file_path_var, state='readonly')
        file_entry.grid(row=0, column=1, sticky=(tk.W, tk.E), padx=(0, 5))
        
        ttk.Button(top_frame, text="Browse...", command=self.open_archive).grid(row=0, column=2)
        
        # Stats bar
        stats_frame = ttk.LabelFrame(main_frame, text="Statistics", padding="5")
        stats_frame.grid(row=1, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
        
        self.stats_text = tk.StringVar(value="No archive loaded")
        ttk.Label(stats_frame, textvariable=self.stats_text).grid(row=0, column=0, sticky=tk.W)
        
        # File list frame
        list_frame = ttk.LabelFrame(main_frame, text="Files in Archive", padding="5")
        list_frame.grid(row=2, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        list_frame.columnconfigure(0, weight=1)
        list_frame.rowconfigure(0, weight=1)
        
        # Create treeview for file list
        columns = ('size', 'compressed', 'ratio', 'flags')
        self.tree = ttk.Treeview(list_frame, columns=columns, show='tree headings')
        
        # Define headings
        self.tree.heading('#0', text='Filename')
        self.tree.heading('size', text='Size')
        self.tree.heading('compressed', text='Compressed')
        self.tree.heading('ratio', text='Ratio')
        self.tree.heading('flags', text='Flags')
        
        # Define column widths
        self.tree.column('#0', width=400)
        self.tree.column('size', width=100, anchor='e')
        self.tree.column('compressed', width=100, anchor='e')
        self.tree.column('ratio', width=80, anchor='e')
        self.tree.column('flags', width=80, anchor='center')
        
        # Scrollbars
        vsb = ttk.Scrollbar(list_frame, orient="vertical", command=self.tree.yview)
        hsb = ttk.Scrollbar(list_frame, orient="horizontal", command=self.tree.xview)
        self.tree.configure(yscrollcommand=vsb.set, xscrollcommand=hsb.set)
        
        self.tree.grid(row=0, column=0, sticky=(tk.N, tk.S, tk.E, tk.W))
        vsb.grid(row=0, column=1, sticky=(tk.N, tk.S))
        hsb.grid(row=1, column=0, sticky=(tk.E, tk.W))
        
        # Context menu for tree
        self.tree_menu = tk.Menu(self.root, tearoff=0)
        self.tree_menu.add_command(label="Extract File...", command=self.extract_selected_file)
        self.tree_menu.add_command(label="View File Info", command=self.show_file_info)
        
        self.tree.bind('<Button-3>', self.show_context_menu)
        self.tree.bind('<Double-1>', self.show_file_info)
        
        # Bottom buttons
        button_frame = ttk.Frame(main_frame)
        button_frame.grid(row=3, column=0, sticky=(tk.W, tk.E), pady=(10, 0))
        
        ttk.Button(button_frame, text="Extract All...", command=self.extract_all_files).pack(side=tk.LEFT, padx=(0, 5))
        ttk.Button(button_frame, text="Refresh", command=self.refresh_file_list).pack(side=tk.LEFT, padx=(0, 5))
        ttk.Button(button_frame, text="Verify Archive", command=self.verify_archive).pack(side=tk.LEFT)
        
        ttk.Button(button_frame, text="Close", command=self.root.quit).pack(side=tk.RIGHT)
    
    def format_size(self, size_bytes):
        """Format size in bytes to human-readable format"""
        for unit in ['B', 'KB', 'MB', 'GB']:
            if size_bytes < 1024.0:
                return f"{size_bytes:.1f} {unit}"
            size_bytes /= 1024.0
        return f"{size_bytes:.1f} TB"
    
    def open_archive(self):
        """Open an MPQ archive"""
        filename = filedialog.askopenfilename(
            title="Select MPQ Archive",
            filetypes=[
                ("MPQ Archives", "*.mpq"),
                ("SC2 Archives", "*.SC2Data *.SC2Map *.SC2Mod"),
                ("All Files", "*.*")
            ]
        )
        
        if not filename:
            return
        
        try:
            # Close current archive if any
            self.close_archive()
            
            # Open new archive
            self.current_archive = MPQArchive(filename)
            self.current_path = filename
            
            # Update UI
            self.file_path_var.set(filename)
            self.refresh_file_list()
            
        except StormLibError as e:
            messagebox.showerror("Error", f"Failed to open archive:\n{str(e)}")
        except Exception as e:
            messagebox.showerror("Error", f"Unexpected error:\n{str(e)}")
    
    def close_archive(self):
        """Close the current archive"""
        if self.current_archive:
            try:
                self.current_archive.close()
            except:
                pass
            self.current_archive = None
            self.current_path = None
        
        self.file_path_var.set("No archive loaded")
        self.stats_text.set("No archive loaded")
        
        # Clear tree
        for item in self.tree.get_children():
            self.tree.delete(item)
    
    def refresh_file_list(self):
        """Refresh the file list from the current archive"""
        if not self.current_archive:
            return
        
        # Clear existing items
        for item in self.tree.get_children():
            self.tree.delete(item)
        
        try:
            # Get all files
            files = self.current_archive.find_files("*")
            
            # Update stats
            total_size = sum(f['size'] for f in files)
            total_compressed = sum(f['compressed_size'] for f in files)
            
            if total_size > 0:
                ratio = (1 - total_compressed / total_size) * 100
            else:
                ratio = 0
            
            self.stats_text.set(
                f"Files: {len(files)} | "
                f"Total Size: {self.format_size(total_size)} | "
                f"Compressed: {self.format_size(total_compressed)} | "
                f"Ratio: {ratio:.1f}%"
            )
            
            # Add files to tree
            for file_info in sorted(files, key=lambda x: x['name']):
                name = file_info['name']
                size = file_info['size']
                compressed = file_info['compressed_size']
                flags = file_info['flags']
                
                # Calculate compression ratio
                if size > 0:
                    file_ratio = (1 - compressed / size) * 100
                else:
                    file_ratio = 0
                
                # Format flags
                flag_str = []
                if flags & 0x00000200:  # MPQ_FILE_COMPRESS
                    flag_str.append("C")
                if flags & 0x00010000:  # MPQ_FILE_ENCRYPTED
                    flag_str.append("E")
                if flags & 0x01000000:  # MPQ_FILE_SINGLE_UNIT
                    flag_str.append("S")
                flag_display = "".join(flag_str) or "-"
                
                self.tree.insert('', 'end', text=name, values=(
                    self.format_size(size),
                    self.format_size(compressed),
                    f"{file_ratio:.1f}%",
                    flag_display
                ))
            
        except Exception as e:
            messagebox.showerror("Error", f"Failed to read archive:\n{str(e)}")
    
    def show_context_menu(self, event):
        """Show context menu for tree item"""
        if not self.current_archive:
            return
        
        item = self.tree.identify_row(event.y)
        if item:
            self.tree.selection_set(item)
            self.tree_menu.post(event.x_root, event.y_root)
    
    def extract_selected_file(self):
        """Extract the selected file"""
        if not self.current_archive:
            return
        
        selection = self.tree.selection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a file to extract")
            return
        
        item = selection[0]
        filename = self.tree.item(item, 'text')
        
        # Ask where to save
        output_path = filedialog.asksaveasfilename(
            title="Save File As",
            initialfile=Path(filename).name,
            defaultextension="*",
            filetypes=[("All Files", "*.*")]
        )
        
        if not output_path:
            return
        
        try:
            self.current_archive.extract_file(filename, output_path)
            messagebox.showinfo("Success", f"File extracted to:\n{output_path}")
        except Exception as e:
            messagebox.showerror("Error", f"Failed to extract file:\n{str(e)}")
    
    def extract_all_files(self):
        """Extract all files from the archive"""
        if not self.current_archive:
            messagebox.showwarning("Warning", "No archive loaded")
            return
        
        # Ask for output directory
        output_dir = filedialog.askdirectory(title="Select Output Directory")
        if not output_dir:
            return
        
        output_dir = Path(output_dir)
        
        try:
            files = self.current_archive.find_files("*")
            total = len(files)
            
            if total == 0:
                messagebox.showinfo("Info", "No files to extract")
                return
            
            # Create progress window
            progress_window = tk.Toplevel(self.root)
            progress_window.title("Extracting Files")
            progress_window.geometry("400x120")
            progress_window.transient(self.root)
            progress_window.grab_set()
            
            ttk.Label(progress_window, text="Extracting files...").pack(pady=(10, 5))
            
            progress_var = tk.StringVar(value="0 / 0")
            ttk.Label(progress_window, textvariable=progress_var).pack()
            
            progress = ttk.Progressbar(progress_window, length=350, mode='determinate', maximum=total)
            progress.pack(pady=10)
            
            extracted = 0
            failed = 0
            
            for i, file_info in enumerate(files, 1):
                filename = file_info['name']
                output_path = output_dir / filename
                output_path.parent.mkdir(parents=True, exist_ok=True)
                
                try:
                    self.current_archive.extract_file(filename, str(output_path))
                    extracted += 1
                except:
                    failed += 1
                
                progress['value'] = i
                progress_var.set(f"{i} / {total}")
                progress_window.update()
            
            progress_window.destroy()
            
            messagebox.showinfo(
                "Extraction Complete",
                f"Extracted: {extracted}\nFailed: {failed}\n\nOutput directory:\n{output_dir}"
            )
            
        except Exception as e:
            messagebox.showerror("Error", f"Extraction failed:\n{str(e)}")
    
    def show_file_info(self, event=None):
        """Show detailed information about selected file"""
        if not self.current_archive:
            return
        
        selection = self.tree.selection()
        if not selection:
            return
        
        item = selection[0]
        filename = self.tree.item(item, 'text')
        values = self.tree.item(item, 'values')
        
        # Create info window
        info_window = tk.Toplevel(self.root)
        info_window.title(f"File Info - {Path(filename).name}")
        info_window.geometry("500x300")
        info_window.transient(self.root)
        
        # Create text widget
        text = scrolledtext.ScrolledText(info_window, wrap=tk.WORD, font=('Courier', 10))
        text.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # Add info
        info_text = f"""File Information
{'='*60}

Filename: {filename}
Size: {values[0]}
Compressed Size: {values[1]}
Compression Ratio: {values[2]}
Flags: {values[3]}

Archive: {self.current_path}
"""
        
        text.insert('1.0', info_text)
        text.config(state=tk.DISABLED)
        
        # Add close button
        ttk.Button(info_window, text="Close", command=info_window.destroy).pack(pady=(0, 10))
    
    def verify_archive(self):
        """Verify the archive integrity"""
        if not self.current_archive:
            messagebox.showwarning("Warning", "No archive loaded")
            return
        
        try:
            result = self.current_archive.verify()
            if result == 0:
                messagebox.showinfo("Verification", "Archive verification successful!\n\nThe archive is valid.")
            else:
                messagebox.showwarning("Verification", f"Archive verification failed!\n\nError code: {result}")
        except Exception as e:
            messagebox.showerror("Error", f"Verification failed:\n{str(e)}")
    
    def show_about(self):
        """Show about dialog"""
        about_text = """MPQ Inspector
Version 1.0.0

A GUI tool for inspecting MPQ archives using PyStorm.

PyStorm: Python bindings for StormLib
StormLib: Created by Ladislav Zezula

Licensed under MIT License
"""
        messagebox.showinfo("About MPQ Inspector", about_text)


def main():
    """Main entry point"""
    root = tk.Tk()
    app = MPQInspectorApp(root)
    root.mainloop()


if __name__ == "__main__":
    main()