#!/usr/bin/env python3
"""
3DE to Nuke Track Converter
A professional tool for converting 3DEqualizer 2D tracks to Nuke Tracker nodes
"""

import argparse
import sys
import json
from pathlib import Path
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass

from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn, TimeElapsedColumn
from rich.prompt import Prompt, Confirm
from rich.layout import Layout
from rich.text import Text
from rich import box
from rich.style import Style

console = Console()

# ASCII Art Banner
BANNER = """
╔═════════════════════════════════════════════════════════════════════════════════════╗
║                                                                                     ║
║ __________.___________    ___ ___    _____  __________  _____ __________________    ║
║ \______   \   \_____  \  /   |   \  /  _  \ \____    / /  _  \\______   \______ \    ║
║  |    |  _/   |/   |   \/    ~    \/  /_\  \  /     / /  /_\  \|       _/|    |  \  ║
║  |    |   \   /    |    \    Y    /    |    \/     /_/    |    \    |   \|    `   \ ║
║  |______  /___\_______  /\___|_  /\____|__  /_______ \____|__  /____|_  /_______  / ║
║         \/            \/       \/         \/        \/       \/       \/        \/  ║
║                                                                                     ║
║ _________ .____    .___  ___________________   ________  .____       _________      ║
║ \_   ___ \|    |   |   | \__    ___/\_____  \  \_____  \ |    |     /   _____/      ║
║ /    \  \/|    |   |   |   |    |    /   |   \  /   |   \|    |     \_____  \       ║
║ \     \___|    |___|   |   |    |   /    |    \/    |    \    |___  /        \      ║
║  \______  /_______ \___|   |____|   \_______  /\_______  /_______ \/_______  /      ║
║         \/        \/                        \/         \/        \/        \/       ║
║                                                                                     ║
║ Convert 2d tracks from 3dequalizer to .nk 2d tracker nodes.                         ║
╚═════════════════════════════════════════════════════════════════════════════════════╝
"""

@dataclass
class TrackPoint:
    """Represents a single track point"""
    frame: int
    x: float
    y: float

@dataclass
class Track:
    """Represents a complete track"""
    name: str
    frame_start: int
    points: List[TrackPoint]
    
    @property
    def frame_end(self) -> int:
        return max(p.frame for p in self.points) if self.points else self.frame_start
    
    @property
    def num_points(self) -> int:
        return len(self.points)

class TrackConverter:
    """Main converter class"""
    
    def __init__(self, config: Dict):
        self.config = config
        self.tracks: List[Track] = []
        
    def parse_3de_file(self, filepath: Path) -> bool:
        """Parse 3DE track file with robust error handling"""
        try:
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
                TimeElapsedColumn(),
                console=console
            ) as progress:
                task = progress.add_task("[cyan]Reading 3DE file...", total=100)
                
                with open(filepath, 'r') as f:
                    lines = f.read().strip().split('\n')
                
                progress.update(task, advance=20)
                
                if not lines:
                    console.print("[red]✗[/red] Error: File is empty", style="bold red")
                    return False
                
                try:
                    track_count = int(lines[0])
                except ValueError:
                    console.print(f"[red]✗[/red] Error: Invalid track count on line 1: '{lines[0]}'", style="bold red")
                    return False
                
                progress.update(task, advance=10, description=f"[cyan]Parsing {track_count} tracks...")
                
                current_line = 1
                parse_increment = 70 / max(track_count, 1)
                
                for i in range(track_count):
                    try:
                        # Parse track number
                        if current_line >= len(lines):
                            raise ValueError(f"Unexpected end of file at track {i+1}")
                        
                        track_num = lines[current_line].strip()
                        current_line += 1
                        
                        # Parse frame start
                        if current_line >= len(lines):
                            raise ValueError(f"Missing frame_start for track {i+1}")
                        frame_start = int(lines[current_line])
                        current_line += 1
                        
                        # Parse number of points
                        if current_line >= len(lines):
                            raise ValueError(f"Missing num_points for track {i+1}")
                        num_points = int(lines[current_line])
                        current_line += 1
                        
                        # Parse points
                        points = []
                        for j in range(num_points):
                            if current_line >= len(lines):
                                raise ValueError(f"Missing point data for track {i+1}, point {j+1}")
                            
                            parts = lines[current_line].strip().split()
                            if len(parts) < 3:
                                raise ValueError(f"Invalid point data at line {current_line+1}: '{lines[current_line]}'")
                            
                            points.append(TrackPoint(
                                frame=int(parts[0]),
                                x=float(parts[1]),
                                y=float(parts[2])
                            ))
                            current_line += 1
                        
                        # Create track with custom name if available, otherwise use track number
                        track_name = self.config.get('name_prefix', '') + f"track_{track_num}"
                        
                        self.tracks.append(Track(
                            name=track_name,
                            frame_start=frame_start,
                            points=points
                        ))
                        
                        progress.update(task, advance=parse_increment)
                        
                    except (ValueError, IndexError) as e:
                        console.print(f"[red]✗[/red] Error parsing track {i+1}: {str(e)}", style="bold red")
                        return False
                
                progress.update(task, completed=100)
            
            return True
            
        except FileNotFoundError:
            console.print(f"[red]✗[/red] File not found: {filepath}", style="bold red")
            return False
        except Exception as e:
            console.print(f"[red]✗[/red] Unexpected error: {str(e)}", style="bold red")
            return False
    
    def transform_coordinates(self, track: Track) -> Track:
        """Transform coordinates based on config settings"""
        if self.config.get('input_space') == 'normalized':
            # Convert normalized (0-1) to pixels
            width = self.config.get('input_width', 2048)
            height = self.config.get('input_height', 1556)
            
            for point in track.points:
                point.x *= width
                point.y *= height
        
        if self.config.get('flip_y', False):
            height = self.config.get('input_height', 1556)
            for point in track.points:
                point.y = height - point.y
        
        return track
    
    def fill_gaps(self, track: Track) -> Track:
        """Fill missing frames in track"""
        if not self.config.get('fill_gaps', False) or not track.points:
            return track
        
        strategy = self.config.get('fill_strategy', 'last')
        filled_points = []
        
        # Sort points by frame
        sorted_points = sorted(track.points, key=lambda p: p.frame)
        
        for i in range(len(sorted_points) - 1):
            current = sorted_points[i]
            next_point = sorted_points[i + 1]
            filled_points.append(current)
            
            # Check for gap
            gap = next_point.frame - current.frame
            if gap > 1:
                for frame in range(current.frame + 1, next_point.frame):
                    if strategy == 'last':
                        filled_points.append(TrackPoint(frame, current.x, current.y))
                    elif strategy == 'linear':
                        # Linear interpolation
                        t = (frame - current.frame) / gap
                        x = current.x + (next_point.x - current.x) * t
                        y = current.y + (next_point.y - current.y) * t
                        filled_points.append(TrackPoint(frame, x, y))
        
        filled_points.append(sorted_points[-1])
        track.points = filled_points
        return track
    
    def generate_curve_data(self, points: List[TrackPoint], axis: str) -> str:
        """Generate Nuke curve data"""
        curve_data = '{curve'
        for point in points:
            value = point.x if axis == 'x' else point.y
            curve_data += f' x{point.frame} {value}'
        curve_data += '}'
        return curve_data
    
    def generate_nuke_file(self, output_path: Path) -> bool:
        """Generate Nuke .nk file"""
        try:
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                console=console
            ) as progress:
                task = progress.add_task("[cyan]Generating Nuke file...", total=len(self.tracks))
                
                # Transform and process tracks
                processed_tracks = []
                for track in self.tracks:
                    track = self.transform_coordinates(track)
                    track = self.fill_gaps(track)
                    processed_tracks.append(track)
                    progress.advance(task)
                
                # Generate track entries
                track_entries = []
                for index, track in enumerate(processed_tracks):
                    track_x = self.generate_curve_data(track.points, 'x')
                    track_y = self.generate_curve_data(track.points, 'y')
                    
                    enabled = '1' if (self.config.get('enable_all', False) or index == 0) else '0'
                    entry = f' {{ {{curve K x1 1}} "{track.name}" {track_x} {track_y} {{curve K x1 0}} {{curve K x1 0}} {enabled} 0 0 {{curve x1 0}} 1 0 -32 -32 32 32 -22 -22 22 22 {{}} {{}}  {{}}  {{}}  {{}}  {{}}  {{}}  {{}}  {{}}  {{}}  {{}}   }}'
                    track_entries.append(entry)
                
                # Get config values
                nuke_version = self.config.get('nuke_version', '15.2 v3')
                format_name = self.config.get('format', '2K_Super_35(full-ap)')
                width = self.config.get('input_width', 2048)
                height = self.config.get('input_height', 1556)
                center_x = width // 2
                center_y = height // 2
                
                # Generate .nk file
                nuke_file = f"""#! C:/Program Files/Nuke15.2v3/nuke-15.2.3.dll -nx
version {nuke_version}
Root {{
inputs 0
name {output_path.name}
format "{width} {height} 0 0 {width} {height} 1 {format_name}"
proxy_type scale
colorManagement OCIO
OCIO_config fn-nuke_studio-config-v1.0.0_aces-v1.3_ocio-v2.1
defaultViewerLUT "OCIO LUTs"
workingSpaceLUT scene_linear
monitorLut "ACES 1.0 - SDR Video (sRGB - Display)"
monitorOutLUT "ACES 1.0 - SDR Video (sRGB - Display)"
}}
Tracker4 {{
inputs 0
tracks {{ {{ 1 31 {len(processed_tracks)} }} 
{{ {{ 5 1 20 enable e 1 }} 
{{ 3 1 75 name name 1 }} 
{{ 2 1 58 track_x track_x 1 }} 
{{ 2 1 58 track_y track_y 1 }} 
{{ 2 1 63 offset_x offset_x 1 }} 
{{ 2 1 63 offset_y offset_y 1 }} 
{{ 4 1 27 T T 1 }} 
{{ 4 1 27 R R 1 }} 
{{ 4 1 27 S S 1 }} 
{{ 2 0 45 error error 1 }} 
{{ 1 1 0 error_min error_min 1 }} 
{{ 1 1 0 error_max error_max 1 }} 
{{ 1 1 0 pattern_x pattern_x 1 }} 
{{ 1 1 0 pattern_y pattern_y 1 }} 
{{ 1 1 0 pattern_r pattern_r 1 }} 
{{ 1 1 0 pattern_t pattern_t 1 }} 
{{ 1 1 0 search_x search_x 1 }} 
{{ 1 1 0 search_y search_y 1 }} 
{{ 1 1 0 search_r search_r 1 }} 
{{ 1 1 0 search_t search_t 1 }} 
{{ 2 1 0 key_track key_track 1 }} 
{{ 2 1 0 key_search_x key_search_x 1 }} 
{{ 2 1 0 key_search_y key_search_y 1 }} 
{{ 2 1 0 key_search_r key_search_r 1 }} 
{{ 2 1 0 key_search_t key_search_t 1 }} 
{{ 2 1 0 key_track_x key_track_x 1 }} 
{{ 2 1 0 key_track_y key_track_y 1 }} 
{{ 2 1 0 key_track_r key_track_r 1 }} 
{{ 2 1 0 key_track_t key_track_t 1 }} 
{{ 2 1 0 key_centre_offset_x key_centre_offset_x 1 }} 
{{ 2 1 0 key_centre_offset_y key_centre_offset_y 1 }} 
}} 
{{ 
{chr(10).join(track_entries)}
}} 
}}
center {{{center_x} {center_y}}}
name Tracker1
xpos -75
ypos -94
}}
"""
                
                with open(output_path, 'w') as f:
                    f.write(nuke_file)
            
            return True
            
        except Exception as e:
            console.print(f"[red]✗[/red] Error generating Nuke file: {str(e)}", style="bold red")
            return False
    
    def display_summary(self):
        """Display conversion summary"""
        if not self.tracks:
            return
        
        # Create summary table
        table = Table(
            title="Track Summary",
            box=box.ROUNDED,
            show_header=True,
            header_style="bold magenta",
            border_style="cyan"
        )
        
        table.add_column("Track", style="cyan", no_wrap=True)
        table.add_column("Points", justify="right", style="green")
        table.add_column("Frame Range", justify="center", style="yellow")
        table.add_column("Status", justify="center")
        
        for track in self.tracks[:10]:  # Show first 10
            status = "[green]✓[/green]" if track.num_points > 0 else "[red]✗[/red]"
            table.add_row(
                track.name,
                str(track.num_points),
                f"{track.frame_start}-{track.frame_end}",
                status
            )
        
        if len(self.tracks) > 10:
            table.add_row("...", "...", "...", "...")
        
        console.print(table)
        
        # Stats panel
        total_points = sum(t.num_points for t in self.tracks)
        stats = f"""
[cyan]Total Tracks:[/cyan] [bold]{len(self.tracks)}[/bold]
[cyan]Total Points:[/cyan] [bold]{total_points}[/bold]
[cyan]Avg Points/Track:[/cyan] [bold]{total_points // len(self.tracks) if self.tracks else 0}[/bold]
"""
        console.print(Panel(stats, title="Statistics", border_style="green", box=box.ROUNDED))

def interactive_mode():
    """Run in interactive mode with prompts"""
    console.print(Panel(BANNER, style="bold cyan", box=box.DOUBLE))
    
    console.print("\n[bold cyan]Welcome to 3DE to Nuke Track Converter![/bold cyan]\n")
    
    # Get input file
    input_file = Prompt.ask(
        "[cyan]Enter 3DE track file path[/cyan] [dim](e.g., tracks.txt from 3DEqualizer)[/dim]",
        default="tracks.txt"
    )
    
    if not Path(input_file).exists():
        console.print(f"[red]✗[/red] File not found: {input_file}", style="bold red")
        return
    
    # Auto-generate output filename from input
    input_path = Path(input_file)
    default_output = input_path.stem + ".nk"
    
    # Get output file
    output_file = Prompt.ask(
        "[cyan]Enter output .nk file path[/cyan]",
        default=default_output
    )
    
    # Resolution
    width = int(Prompt.ask("[cyan]Input resolution width[/cyan]", default="2048"))
    height = int(Prompt.ask("[cyan]Input resolution height[/cyan]", default="1556"))
    
    # Coordinate space - removed confusing option, using auto-detection
    console.print("\n[dim]Note: Coordinate space auto-detected as pixels (standard 3DE export)[/dim]")
    coord_space = "pixels"
    
    # Fill gaps
    fill_gaps = Confirm.ask("[cyan]Fill missing frames?[/cyan]", default=False)
    fill_strategy = "last"
    if fill_gaps:
        fill_strategy = Prompt.ask(
            "[cyan]Fill strategy[/cyan]",
            choices=["last", "linear"],
            default="last"
        )
    
    # Enable all tracks
    enable_all = Confirm.ask("[cyan]Enable all tracks?[/cyan]", default=False)
    
    # Build config
    config = {
        'input_width': width,
        'input_height': height,
        'input_space': coord_space,
        'fill_gaps': fill_gaps,
        'fill_strategy': fill_strategy,
        'enable_all': enable_all,
        'flip_y': False,
        'name_prefix': '',
        'nuke_version': '15.2 v3',
        'format': '2K_Super_35(full-ap)'
    }
    
    # Process
    converter = TrackConverter(config)
    
    console.print("\n")
    if converter.parse_3de_file(Path(input_file)):
        console.print("[green]✓[/green] Successfully parsed 3DE file\n", style="bold green")
        converter.display_summary()
        
        console.print("\n")
        if converter.generate_nuke_file(Path(output_file)):
            console.print(f"\n[green]✓[/green] Successfully created: [bold]{output_file}[/bold]", style="bold green")
            
            # Final summary
            summary = Panel(
                f"[green]Conversion Complete![/green]\n\n"
                f"[cyan]Input:[/cyan] {input_file}\n"
                f"[cyan]Output:[/cyan] {output_file}\n"
                f"[cyan]Tracks:[/cyan] {len(converter.tracks)}\n"
                f"[cyan]Total Points:[/cyan] {sum(t.num_points for t in converter.tracks)}",
                title="✓ Success",
                border_style="green",
                box=box.DOUBLE
            )
            console.print(summary)

def cli_mode(args):
    """Run in CLI mode with arguments"""
    config = {
        'input_width': args.width,
        'input_height': args.height,
        'input_space': args.input_space,
        'fill_gaps': args.fill_gaps,
        'fill_strategy': args.fill_strategy,
        'enable_all': args.enable_all,
        'flip_y': args.flip_y,
        'name_prefix': args.name_prefix,
        'nuke_version': args.nuke_version,
        'format': args.format
    }
    
    if not args.quiet:
        console.print(Panel(BANNER, style="bold cyan", box=box.DOUBLE))
    
    converter = TrackConverter(config)
    
    if converter.parse_3de_file(Path(args.input)):
        if not args.quiet:
            console.print("[green]✓[/green] Successfully parsed 3DE file\n", style="bold green")
            if args.verbose:
                converter.display_summary()
        
        if args.dry_run:
            console.print("[yellow]Dry run - no file written[/yellow]")
            return
        
        if converter.generate_nuke_file(Path(args.output)):
            if not args.quiet:
                console.print(f"\n[green]✓[/green] Successfully created: [bold]{args.output}[/bold]", style="bold green")
        else:
            sys.exit(1)
    else:
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(
        description='Convert 3DEqualizer 2D tracks to Nuke Tracker nodes',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    parser.add_argument('-i', '--input', help='Input 3DE track file')
    parser.add_argument('-o', '--output', help='Output Nuke .nk file')
    parser.add_argument('-w', '--width', type=int, default=2048, help='Input resolution width (default: 2048)')
    parser.add_argument('--height', type=int, default=1556, help='Input resolution height (default: 1556)')
    parser.add_argument('--input-space', choices=['pixels', 'normalized'], default='pixels', 
                       help='Input coordinate space (default: pixels)')
    parser.add_argument('--fill-gaps', action='store_true', help='Fill missing frames')
    parser.add_argument('--fill-strategy', choices=['last', 'linear'], default='last',
                       help='Gap filling strategy (default: last)')
    parser.add_argument('--flip-y', action='store_true', help='Flip Y coordinates')
    parser.add_argument('--enable-all', action='store_true', help='Enable all tracks (default: only first)')
    parser.add_argument('--name-prefix', default='', help='Prefix for track names')
    parser.add_argument('--nuke-version', default='15.2 v3', help='Nuke version (default: 15.2 v3)')
    parser.add_argument('--format', default='2K_Super_35(full-ap)', help='Nuke format name')
    parser.add_argument('-v', '--verbose', action='store_true', help='Verbose output')
    parser.add_argument('-q', '--quiet', action='store_true', help='Quiet mode')
    parser.add_argument('--dry-run', action='store_true', help='Parse only, don\'t write file')
    
    # If no arguments, run interactive mode
    if len(sys.argv) == 1:
        interactive_mode()
    else:
        args = parser.parse_args()
        if not args.input or not args.output:
            parser.error("--input and --output are required in CLI mode")
        cli_mode(args)

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        console.print("\n[yellow]Cancelled by user[/yellow]")
        sys.exit(0)
    except Exception as e:
        console.print(f"\n[red]Fatal error: {str(e)}[/red]", style="bold red")
        sys.exit(1)