import os
import subprocess
from pathlib import Path
import json
import logging
from datetime import datetime

# ANSI color codes
class Colors:
    PURPLE = '\033[95m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    ENDC = '\033[0m'

# Set up logging
log_dir = "logs"
os.makedirs(log_dir, exist_ok=True)
log_file = os.path.join(log_dir, f"repair_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
logging.basicConfig(filename=log_file, level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def check_file_health(input_file):
    """Check if MP4 file has issues using ffmpeg"""
    print(f"\n{Colors.BLUE}Analyzing: {input_file}{Colors.ENDC}")
    logging.info(f"Analyzing file: {input_file}")
    
    cmd = [
        'ffmpeg',
        '-v', 'error',
        '-i', str(input_file),
        '-f', 'null',
        '-'
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    if result.stderr:
        print(f"{Colors.RED}Issues found:{Colors.ENDC}")
        for line in result.stderr.split('\n')[:10]:  # Show first 10 errors
            if line.strip():
                print(f"  {Colors.YELLOW}{line}{Colors.ENDC}")
        logging.warning(f"File has errors: {result.stderr[:500]}")
        return False
    else:
        print(f"{Colors.GREEN}No errors detected{Colors.ENDC}")
        logging.info("File appears healthy")
        return True

def get_file_info(input_file):
    """Get basic file info"""
    cmd = [
        'ffprobe',
        '-v', 'error',
        '-show_entries', 'format=duration,size:stream=codec_name,width,height',
        '-of', 'json',
        input_file
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    try:
        return json.loads(result.stdout)
    except:
        return None

def repair_method_1_remux(input_file, output_file):
    """Method 1: Simple remux (copy streams to new container)"""
    print(f"\n{Colors.BLUE}Trying Method 1: Remux (copy streams){Colors.ENDC}")
    logging.info("Attempting repair method 1: Remux")
    
    cmd = [
        'ffmpeg',
        '-v', 'warning',
        '-i', str(input_file),
        '-c', 'copy',
        '-map', '0',
        '-y',
        str(output_file)
    ]
    
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        if result.returncode == 0 and Path(output_file).exists():
            print(f"{Colors.GREEN}Method 1 successful!{Colors.ENDC}")
            logging.info("Method 1 successful")
            return True
        else:
            print(f"{Colors.RED}Method 1 failed{Colors.ENDC}")
            logging.warning(f"Method 1 failed: {result.stderr[:200]}")
            return False
    except Exception as e:
        print(f"{Colors.RED}Method 1 error: {e}{Colors.ENDC}")
        logging.error(f"Method 1 error: {e}")
        return False

def repair_method_2_error_concealment(input_file, output_file):
    """Method 2: Remux with error concealment"""
    print(f"\n{Colors.BLUE}Trying Method 2: Error concealment{Colors.ENDC}")
    logging.info("Attempting repair method 2: Error concealment")
    
    cmd = [
        'ffmpeg',
        '-v', 'warning',
        '-err_detect', 'ignore_err',
        '-i', str(input_file),
        '-c', 'copy',
        '-map', '0',
        '-fflags', '+genpts+igndts',
        '-y',
        str(output_file)
    ]
    
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        if result.returncode == 0 and Path(output_file).exists():
            print(f"{Colors.GREEN}Method 2 successful!{Colors.ENDC}")
            logging.info("Method 2 successful")
            return True
        else:
            print(f"{Colors.RED}Method 2 failed{Colors.ENDC}")
            logging.warning(f"Method 2 failed: {result.stderr[:200]}")
            return False
    except Exception as e:
        print(f"{Colors.RED}Method 2 error: {e}{Colors.ENDC}")
        logging.error(f"Method 2 error: {e}")
        return False

def repair_method_3_force_keyframes(input_file, output_file):
    """Method 3: Force keyframe processing"""
    print(f"\n{Colors.BLUE}Trying Method 3: Force keyframe processing{Colors.ENDC}")
    logging.info("Attempting repair method 3: Force keyframe processing")
    
    cmd = [
        'ffmpeg',
        '-v', 'warning',
        '-fflags', '+discardcorrupt+genpts',
        '-i', str(input_file),
        '-c', 'copy',
        '-map', '0',
        '-avoid_negative_ts', 'make_zero',
        '-y',
        str(output_file)
    ]
    
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        if result.returncode == 0 and Path(output_file).exists():
            print(f"{Colors.GREEN}Method 3 successful!{Colors.ENDC}")
            logging.info("Method 3 successful")
            return True
        else:
            print(f"{Colors.RED}Method 3 failed{Colors.ENDC}")
            logging.warning(f"Method 3 failed: {result.stderr[:200]}")
            return False
    except Exception as e:
        print(f"{Colors.RED}Method 3 error: {e}{Colors.ENDC}")
        logging.error(f"Method 3 error: {e}")
        return False

def repair_file(input_file, output_dir):
    """Attempt to repair a damaged MP4 file"""
    input_path = Path(input_file)
    output_path = Path(output_dir) / f"{input_path.stem}_repaired{input_path.suffix}"
    
    print(f"\n{Colors.PURPLE}{'='*60}{Colors.ENDC}")
    print(f"{Colors.PURPLE}Processing: {input_path.name}{Colors.ENDC}")
    print(f"{Colors.PURPLE}{'='*60}{Colors.ENDC}")
    
    # Skip if already repaired
    if output_path.exists():
        print(f"{Colors.YELLOW}Already repaired: {output_path.name}{Colors.ENDC}")
        logging.info(f"Skipping {input_path} - already repaired")
        return
    
    # Check file health first
    is_healthy = check_file_health(str(input_path))
    
    if is_healthy:
        print(f"{Colors.GREEN}File appears healthy, no repair needed{Colors.ENDC}")
        logging.info("File is healthy, skipping")
        return
    
    # Try repair methods in sequence
    methods = [
        repair_method_1_remux,
        repair_method_2_error_concealment,
        repair_method_3_force_keyframes
    ]
    
    for method in methods:
        if method(str(input_path), str(output_path)):
            # Verify the repaired file
            print(f"\n{Colors.BLUE}Verifying repaired file...{Colors.ENDC}")
            if check_file_health(str(output_path)):
                print(f"{Colors.GREEN}✓ Repair successful and verified!{Colors.ENDC}")
                logging.info(f"Successfully repaired: {output_path}")
                
                # Show size comparison
                original_size = input_path.stat().st_size / (1024*1024)
                repaired_size = output_path.stat().st_size / (1024*1024)
                print(f"Original: {original_size:.2f} MB → Repaired: {repaired_size:.2f} MB")
                logging.info(f"Size: {original_size:.2f} MB → {repaired_size:.2f} MB")
                return
            else:
                print(f"{Colors.YELLOW}Repair completed but file still has issues{Colors.ENDC}")
                logging.warning("Repaired file still has issues")
                return
    
    print(f"{Colors.RED}All repair methods failed{Colors.ENDC}")
    logging.error(f"Failed to repair: {input_path}")

if __name__ == "__main__":
    input_dir = "input"
    output_dir = "output"
    os.makedirs(output_dir, exist_ok=True)
    
    # Get list of MP4 files
    files = [f for f in os.listdir(input_dir) if f.endswith(('.mp4', '.DVR.mp4'))]
    total_files = len(files)
    
    if total_files == 0:
        print(f"{Colors.YELLOW}No MP4 files found in input directory{Colors.ENDC}")
        logging.info("No files to process")
    else:
        print(f"\n{Colors.BLUE}Found {total_files} MP4 files to check{Colors.ENDC}")
        logging.info(f"Found {total_files} files to process")
        
        for i, file in enumerate(files, 1):
            input_file = os.path.join(input_dir, file)
            print(f"\n{Colors.BLUE}[{i}/{total_files}]{Colors.ENDC}")
            repair_file(input_file, output_dir)
        
        print(f"\n{Colors.GREEN}{'='*60}{Colors.ENDC}")
        print(f"{Colors.GREEN}Repair process complete!{Colors.ENDC}")
        print(f"{Colors.GREEN}Check the output directory for repaired files{Colors.ENDC}")
        print(f"{Colors.GREEN}Log file: {log_file}{Colors.ENDC}")
        print(f"{Colors.GREEN}{'='*60}{Colors.ENDC}")