"""Batch conversion utilities"""
import json
import yaml
from pathlib import Path
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn

from alloy.exceptions import ConfigError, UnsupportedModelError
from alloy.utils.errors import get_config_suggestions

console = Console()

def parse_batch_file(batch_file):
    """
    Parse batch conversion file (JSON or YAML).
    
    Expected format:
    [
      {
        "model": "black-forest-labs/FLUX.1-schnell",
        "type": "flux",
        "output_dir": "converted_models/flux_schnell",
        "quantization": "int4"
      },
      ...
    ]
    """
    path = Path(batch_file)
    if not path.exists():
        raise ConfigError(
            f"Batch file not found: {batch_file}",
            config_file=str(path),
            suggestions=["Verify the file path is correct", "Check that the file exists"],
        )
    
    content = path.read_text()
    
    # Try JSON first
    if path.suffix in ['.json', '.jsonl']:
        try:
            data = json.loads(content)
        except json.JSONDecodeError as e:
            raise ConfigError(
                f"Invalid JSON: {e}",
                config_file=str(path),
                suggestions=get_config_suggestions(),
            )
    # Try YAML
    elif path.suffix in ['.yaml', '.yml']:
        try:
            data = yaml.safe_load(content)
        except yaml.YAMLError as e:
            raise ConfigError(
                f"Invalid YAML: {e}",
                config_file=str(path),
                suggestions=get_config_suggestions(),
            )
    else:
        # Try both
        try:
            data = json.loads(content)
        except Exception:
            try:
                data = yaml.safe_load(content)
            except Exception:
                raise ConfigError(
                    "File must be JSON or YAML format",
                    config_file=str(path),
                    suggestions=get_config_suggestions(),
                )

    if not isinstance(data, list):
        raise ConfigError(
            "Batch file must contain a list of model configs",
            config_file=str(path),
            suggestions=["Ensure the file contains a JSON/YAML array of model configs"],
        )
    
    return data


def validate_batch_config(config):
    """Validate a single batch config entry"""
    required = ['model', 'type']
    missing = [field for field in required if field not in config]
    if missing:
        raise ConfigError(
            "Missing required fields",
            missing_fields=missing,
            suggestions=get_config_suggestions(missing_fields=missing),
        )
    
    # Set defaults
    config.setdefault('output_dir', f"converted_models/{config['model'].split('/')[-1]}")
    config.setdefault('quantization', 'float16')
    
    return config


def run_batch_conversion(batch_file, dry_run=False, parallel=False):
    """
    Run batch conversion from file.
    
    Args:
        batch_file: Path to JSON/YAML file with model configs
        dry_run: If True, only print what would be converted
        parallel: If True, run conversions in parallel (experimental)
    """
    console.print(f"[cyan]Loading batch file:[/cyan] {batch_file}")
    
    try:
        configs = parse_batch_file(batch_file)
    except Exception as e:
        console.print(f"[red]Error parsing batch file:[/red] {e}")
        return False
    
    console.print(f"[green]Found {len(configs)} model(s) to convert[/green]")
    
    # Validate all configs first
    validated_configs = []
    for i, config in enumerate(configs, 1):
        try:
            validated = validate_batch_config(config)
            validated_configs.append(validated)
            console.print(f"  [{i}] {validated['model']} → {validated['output_dir']}")
        except Exception as e:
            console.print(f"[red]  [{i}] Invalid config:[/red] {e}")
            return False
    
    if dry_run:
        console.print("\n[yellow]Dry run - no conversions performed[/yellow]")
        return True
    
    # Run conversions
    if parallel:
        console.print("\n[yellow]⚠ Parallel conversion is experimental and may cause high memory usage[/yellow]")
        # TODO: Implement parallel conversion with multiprocessing
        console.print("[red]Parallel mode not yet implemented, falling back to sequential[/red]")
    
    # Sequential conversion
    console.print(f"\n[cyan]Starting sequential conversion...[/cyan]")
    
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
        TimeRemainingColumn(),
    ) as progress:
        overall = progress.add_task("[cyan]Overall Progress", total=len(validated_configs))
        
        success_count = 0
        failed = []
        
        for i, config in enumerate(validated_configs, 1):
            model = config['model']
            progress.update(overall, description=f"Converting {i}/{len(validated_configs)}: {model}")
            
            try:
                # Import here to avoid circular imports
                from alloy.converters.base import SDConverter
                from alloy.converters.flux import FluxConverter
                from alloy.converters.ltx import LTXConverter
                from alloy.converters.wan import WanConverter
                from alloy.converters.hunyuan import HunyuanConverter
                from alloy.converters.lumina import LuminaConverter
                
                # Create appropriate converter
                converter_map = {
                    'sd': SDConverter,
                    'flux': FluxConverter,
                    'ltx': LTXConverter,
                    'wan': WanConverter,
                    'hunyuan': HunyuanConverter,
                    'lumina': LuminaConverter,
                }
                
                converter_class = converter_map.get(config['type'])
                if not converter_class:
                    raise UnsupportedModelError(
                        f"Unknown model type: {config['type']}",
                        model_type=config['type'],
                        supported_types=list(converter_map.keys())
                    )
                
                converter = converter_class(
                    config['model'],
                    config['output_dir'],
                    config['quantization']
                )
                
                console.print(f"\n[bold cyan]Converting {model}...[/bold cyan]")
                converter.convert()
                console.print(f"[green]✓ Completed: {model}[/green]")
                success_count += 1
                
            except Exception as e:
                console.print(f"[red]✗ Failed: {model}[/red]")
                console.print(f"[red]  Error: {e}[/red]")
                failed.append((model, str(e)))
            
            progress.update(overall, advance=1)
    
    # Summary
    console.print(f"\n[bold]Batch Conversion Summary:[/bold]")
    console.print(f"  [green]Success:[/green] {success_count}/{len(validated_configs)}")
    if failed:
        console.print(f"  [red]Failed:[/red] {len(failed)}")
        for model, error in failed:
            console.print(f"    - {model}: {error}")
    
    return len(failed) == 0
