#!/usr/bin/env python3
r"""
Smart VMAF Monitor & Dashboard
------------------------------
The main interface for the VMAF Optimizer.
Provides a TUI (Text User Interface) to visualize encoding progress.

Usage:
    python smart_monitor.py --tv-dir "Z:\tv" --content-dir "Z:\content" --jobs 4 [--monitor]
"""

import os
import time
import argparse
import sys
import threading
import queue
import json
import re
from pathlib import Path

# Import the engine
import smart_gpu_encoder as encoder
import vmaf_common as common

# UI Library
try:
    from rich.console import Console
    from rich.live import Live
    from rich.table import Table
    from rich.layout import Layout
    from rich.panel import Panel
    from rich.text import Text
    HAS_RICH = True
except ImportError:
    HAS_RICH = False
    print("Warning: 'rich' library not found. Running in basic mode.")

# Watchdog
try:
    from watchdog.observers import Observer
    from watchdog.events import FileSystemEventHandler
    HAS_WATCHDOG = True
except ImportError:
    HAS_WATCHDOG = False

# --- Caching ---
CACHE_FILE = Path("library_cache.json")

def load_cache():
    if CACHE_FILE.exists():
        try:
            with open(CACHE_FILE, "r") as f:
                return set(json.load(f))
        except: pass
    return set()

def save_cache(files):
    try:
        with open(CACHE_FILE, "w") as f:
            # Convert paths to strings for JSON
            json.dump([str(f) for f in files], f)
    except: pass

def fast_scan(path):
    """Recursive scan using scandir (faster than pathlib)"""
    files = []
    try:
        if not os.path.exists(path): return []
        for entry in os.scandir(path):
            if entry.is_dir():
                files.extend(fast_scan(entry.path))
            elif entry.is_file():
                if entry.name.lower().endswith(('.mkv', '.mp4')):
                    if "_enc" not in entry.name and "_ref" not in entry.name:
                        files.append(entry.path)
    except:
        pass
    return files

# --- UI State ---
class Dashboard:
    def __init__(self, num_workers, log_dir=None):
        self.num_workers = num_workers
        self.worker_status = {i: {"file": "Idle", "action": "Waiting", "progress": 0, "speed": "", "color": "dim"} for i in range(num_workers)}
        self.stats = {"processed": 0, "skipped": 0, "failed": 0, "rejected": 0, "savings_gb": 0.0}
        self.recent_completed = []
        self.lock = threading.Lock()
        self.log_dir = log_dir
        self.activity_log_file = None
        
        # Open activity log file for streaming
        if log_dir:
            try:
                log_path = Path(log_dir) / "activity.log"
                self.activity_log_file = open(log_path, "a", encoding="utf-8")
            except Exception as e:
                print(f"[Warning] Could not open activity log: {e}")
        
    def format_filename(self, filename):
        # Clean Sonarr format: {Series} - S{s}E{e} - {Title} {Quality}
        # Goal: Series S01E01 [Quality]
        try:
            # Match S01E01
            match = re.search(r"(.*?) - (S\d+E\d+) - .*? ((?:Bluray|WebDL|Remux|2160p|1080p|Proper).*)\.mkv", filename, re.IGNORECASE)
            if match:
                series = match.group(1)[:15] # Truncate series name
                s_e = match.group(2)
                quality = match.group(3).split()[0] # Just take first part of quality (Bluray-2160p)
                return f"{series} {s_e} [{quality}]"
            
            # Fallback for simpler names
            if len(filename) > 35: 
                return filename[:15] + "..." + filename[-15:]
            return filename
        except:
            return filename

    def update_worker(self, worker_id, file, action, color="blue"):
        with self.lock:
            display_name = self.format_filename(file)
            
            progress = 0
            speed = ""
            
            # Parse rich status: "Encoding AV1 | 45.2% | 2.3x"
            if "|" in action:
                parts = action.split("|")
                action_text = parts[0].strip()
                for p in parts[1:]:
                    p = p.strip()
                    if "%" in p:
                        try: progress = float(p.replace("%", ""))
                        except: pass
                    elif "x" in p or "MB/s" in p:
                        speed = p
                action = action_text
            
            self.worker_status[worker_id] = {
                "file": display_name, 
                "action": action, 
                "progress": progress, 
                "speed": speed, 
                "color": color
            }
            
    def add_log(self, message):
        with self.lock:
            ts = time.strftime("%Y-%m-%d %H:%M:%S")
            ts_short = time.strftime("%H:%M:%S")
            log_entry = f"[{ts_short}] {message}"
            
            self.recent_completed.insert(0, log_entry)
            if len(self.recent_completed) > 12:
                self.recent_completed.pop()
            
            # Stream to activity log file
            if self.activity_log_file:
                try:
                    self.activity_log_file.write(f"[{ts}] {message}\n")
                    self.activity_log_file.flush()
                except:
                    pass

    def update_stats(self, key, val=1):
        with self.lock:
            if key == "savings_gb": self.stats[key] += val
            else: self.stats[key] += val

    def get_renderable(self):
        if not HAS_RICH: return ""
        
        layout = Layout()
        layout.split_column(
            Layout(name="header", size=3),
            Layout(name="workers", size=self.num_workers + 4),
            Layout(name="stats", size=3),
            Layout(name="logs", ratio=1)
        )
        
        layout["header"].update(Panel(Text("Smart GPU Encoder (AMD AMF + VMAF)", justify="center", style="bold cyan")))
        
        # Workers Table
        table = Table(box=None, expand=True, padding=(0, 1))
        table.add_column("ID", width=3, style="dim", no_wrap=True)
        table.add_column("File", ratio=6, no_wrap=True)
        table.add_column("Action", ratio=3, no_wrap=True)
        table.add_column("Progress", width=20, no_wrap=True)
        table.add_column("Speed", width=12, no_wrap=True)
        
        for i in range(self.num_workers):
            ws = self.worker_status[i]
            
            pct = ws["progress"]
            # Rich Bar
            if pct > 0:
                bar_len = 12
                filled = int(bar_len * (pct / 100))
                bar_str = "━" * filled + " " * (bar_len - filled)
                prog_render = Text(f"{bar_str} {pct:.1f}%", style="green")
            else:
                prog_render = Text("")
                
            table.add_row(
                str(i+1), 
                Text(ws["file"], style="white"), 
                Text(ws["action"], style=ws["color"]), 
                prog_render,
                Text(ws["speed"], style="yellow")
            )
            
        layout["workers"].update(Panel(table, title="Active Workers"))
        
        stat_str = f"Processed: [green]{self.stats['processed']}[/] | Skipped: [yellow]{self.stats['skipped']}[/] | Rejected: [magenta]{self.stats['rejected']}[/] | Failed: [red]{self.stats['failed']}[/] | Saved: [bold green]{self.stats['savings_gb']:.2f} GB[/]"
        layout["stats"].update(Panel(Text.from_markup(stat_str, justify="center")))
        
        layout["logs"].update(Panel("\n".join(self.recent_completed), title="Activity Log"))
        
        return layout

# --- Worker Bridge ---
def worker_wrapper(worker_id, file_path, category, lock_dir, log_dir, encoders, dashboard):
    def status_callback(w_id, fname, msg, color):
        # Handle Stats Signal
        if msg.startswith("STATS:"):
            parts = msg.split(":")
            if parts[1] == "SAVED":
                try:
                    bytes_saved = float(parts[2])
                    dashboard.update_stats("savings_gb", bytes_saved / (1024**3))
                    dashboard.update_stats("processed")
                except: pass
            return

        # Update Live Table
        dashboard.update_worker(worker_id, fname, msg, color)
        
        # Check for Completion Events to update History
        # Match keywords from smart_gpu_encoder.py
        if "Done" in msg: 
            dashboard.add_log(f"[Success] {dashboard.format_filename(fname)}")
        elif "Skipped" in msg or "Skipping" in msg:
            dashboard.add_log(f"[Skip] {dashboard.format_filename(fname)}")
            dashboard.update_stats("skipped")
        elif "Rejected" in msg:
            dashboard.add_log(f"[Reject] {dashboard.format_filename(fname)}")
            dashboard.update_stats("rejected")
            # Rejected doesn't count as failed, just processed (or ignored stats)
        elif "Failed" in msg or "Error" in msg:
            dashboard.add_log(f"[Fail] {dashboard.format_filename(fname)}")
            dashboard.update_stats("failed")
            
    try:
        encoder.process_file(file_path, category, lock_dir, log_dir, encoders, worker_id, status_callback)
    except Exception as e:
        dashboard.add_log(f"Error in worker {worker_id}: {str(e)[:30]}")
        dashboard.update_stats("failed")
    finally:
        dashboard.update_worker(worker_id, "Idle", "Waiting", "dim")

# --- Monitor ---
if HAS_WATCHDOG:
    class WatcherHandler(FileSystemEventHandler):
        def __init__(self, queue):
            self.queue = queue
        def on_created(self, event):
            if not event.is_directory and event.src_path.lower().endswith(('.mkv', '.mp4')):
                self.queue.put(Path(event.src_path))
        def on_moved(self, event):
            if not event.is_directory and event.dest_path.lower().endswith(('.mkv', '.mp4')):
                self.queue.put(Path(event.dest_path))
else:
    # Dummy class when watchdog not available
    class WatcherHandler:
        def __init__(self, queue):
            self.queue = queue

# --- Main ---
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--tv-dir", default=common.DEFAULT_CONFIG["tv_dir"])
    parser.add_argument("--content-dir", default=common.DEFAULT_CONFIG["content_dir"])
    parser.add_argument("--jobs", type=int, default=4)
    parser.add_argument("--monitor", action="store_true")
    parser.add_argument("--cpu-only", action="store_true", help="Force software encoding")
    parser.add_argument("--temp-dir", help="Override local temp directory")
    parser.add_argument("--skip-until", help="Skip files until this substring is found in filename")
    parser.add_argument("--av1-encoder", choices=["hw", "sw", "off"], default="hw", help="AV1 encoder: hw (hardware), sw (software), off (disable)")
    parser.add_argument("--hevc-encoder", choices=["hw", "sw", "off"], default="hw", help="HEVC encoder: hw (hardware), sw (software), off (disable)")
    args = parser.parse_args()
    
    # Setup
    lock_dir, log_dir = common.get_base_paths(args)
    encoder.TEMP_DIR = common.get_temp_dir(args)
    
    # Detect HW
    encoders = common.detect_hardware_encoder(args)
    
    # UI
    dashboard = Dashboard(args.jobs, log_dir=log_dir)
    dashboard.add_log(f"Logs: {log_dir}")
    
    # Work Queue
    work_queue = queue.Queue()
    
    # Background Scanner
    def background_scanner():
        time.sleep(2) # Let UI start
        dashboard.add_log("Starting background scan...")
        
        # Skip-until logic
        skip_until = args.skip_until
        skipping = bool(skip_until)  # True until we find the match
        skipped_count = 0
        
        if skip_until:
            dashboard.add_log(f"Skip mode: Looking for '{skip_until}'")
        
        def should_skip(filepath):
            nonlocal skipping, skipped_count
            if not skipping:
                return False
            # Check if skip_until substring is in the filepath
            if skip_until.lower() in str(filepath).lower():
                skipping = False  # Found it, stop skipping
                dashboard.add_log(f"Match found: {Path(filepath).name}")
                dashboard.add_log(f"Resuming processing from here")
                return False
            skipped_count += 1
            return True
        
        # Load Cache first
        cached_files = load_cache()
        if cached_files:
            dashboard.add_log(f"Loaded {len(cached_files)} files from cache.")
            for f in cached_files:
                if should_skip(f):
                    continue
                p = Path(f)
                cat = "tv_shows" if str(args.tv_dir) in str(p) else "content"
                work_queue.put((p, cat))
        
        # Real Scan
        all_files = []
        for d, cat in [(args.tv_dir, "tv_shows"), (args.content_dir, "content")]:
            found = fast_scan(d)
            for f in found:
                all_files.append(f)
                # Only add if NOT in cache
                if str(f) not in cached_files:
                    if should_skip(f):
                        continue
                    work_queue.put((Path(f), cat))
        
        if skipped_count > 0:
            dashboard.add_log(f"Skipped {skipped_count} files (--skip-until)")
                    
        dashboard.add_log(f"Scan complete. Total: {len(all_files)}")
        save_cache(all_files)

    # Start Scanner Thread
    scan_thread = threading.Thread(target=background_scanner, daemon=True)
    scan_thread.start()
    
    # Thread Pool for Workers
    threads = []
    
    def worker_loop(w_id):
        while not encoder.shutdown_requested:
            try:
                item = work_queue.get(timeout=1)
                file_path, category = item
                worker_wrapper(w_id, file_path, category, lock_dir, log_dir, encoders, dashboard)
                work_queue.task_done()
            except queue.Empty:
                # If batch mode and scan is done and queue is empty, exit
                if not args.monitor and not scan_thread.is_alive() and work_queue.empty():
                    time.sleep(2) # Grace period
                    if work_queue.empty():
                        return # Exit thread
                continue
            except Exception as e:
                dashboard.add_log(f"Worker {w_id} crashed: {e}")

    for i in range(args.jobs):
        t = threading.Thread(target=worker_loop, args=(i,), daemon=True)
        t.start()
        threads.append(t)

    # UI Loop
    try:
        if HAS_RICH:
            with Live(dashboard.get_renderable(), refresh_per_second=4) as live:
                while not encoder.shutdown_requested:
                    live.update(dashboard.get_renderable())
                    time.sleep(0.25)
                    
                    # Exit condition
                    if not args.monitor and not scan_thread.is_alive() and work_queue.unfinished_tasks == 0:
                        # Check if threads are actually dead
                        if all(not t.is_alive() for t in threads):
                            break
        else:
            while not encoder.shutdown_requested:
                time.sleep(1)
                if not args.monitor and not scan_thread.is_alive() and work_queue.empty(): break
                
    except KeyboardInterrupt:
        encoder.shutdown_requested = True
        print("\nStopping...")
    
    print("\nDone.")

if __name__ == "__main__":
    main()
