"""3Di encoding command."""

from pathlib import Path
from typing import Optional
import traceback

try:
    import typer
except ImportError:
    typer = None


def encode3di_command(
    input: Path = typer.Option(
        ...,
        "--input",
        "-i",
        help="Input JSON file with protein records",
        exists=True,
        dir_okay=False,
    ),
    output: Path = typer.Option(
        ...,
        "--output",
        "-o",
        help="Output JSON file with 3Di records",
    ),
    model: str = typer.Option(
        "Rostlab/ProstT5_fp16",
        "--model",
        "-m",
        help="ProstT5 model name",
    ),
    device: Optional[str] = typer.Option(
        None,
        "--device",
        "-d",
        help="Device to use (auto/cuda/mps/cpu). Ignored if --multi-gpu is set.",
    ),
    encoding_size: int = typer.Option(
        10000,
        "--encoding-size",
        "-e",
        help="Encoding size (approximates to amino acids)",
    ),
    multi_gpu: bool = typer.Option(
        False,
        "--multi-gpu",
        help="Use multi-GPU parallel encoding when available",
    ),
    gpu_ids: Optional[str] = typer.Option(
        None,
        "--gpu-ids",
        help="Comma-separated list of GPU IDs to use (e.g., '0,1,2'). "
             "If not specified, auto-discovers available GPUs.",
    ),
) -> None:
    """Encode proteins to 3Di structural tokens.

    Uses ProstT5 model to predict 3Di structural alphabet tokens
    directly from amino acid sequences.
    
    Multi-GPU encoding can significantly speed up encoding by distributing
    batches across multiple GPUs. Use --multi-gpu to enable, and optionally
    specify --gpu-ids to select specific GPUs.
    """
    try:
        from ...encode3di.prostt5 import ProstT5ThreeDiEncoder
        from ...io.jsonio import read_json, write_json
        from ...orf.types import OrfRecord
        from ...translate.translator import ProteinRecord

        # Parse GPU IDs if provided
        parsed_gpu_ids = None
        if gpu_ids:
            try:
                parsed_gpu_ids = [int(x.strip()) for x in gpu_ids.split(",")]
            except ValueError:
                typer.echo(f"Error: Invalid GPU IDs format: {gpu_ids}", err=True)
                typer.echo("Expected comma-separated integers, e.g., '0,1,2'", err=True)
                raise typer.Exit(2)

        typer.echo(f"Reading proteins from: {input}")
        protein_data = read_json(input)

        # Reconstruct ProteinRecord objects
        if isinstance(protein_data, list):
            proteins = []
            for p in protein_data:
                orf = OrfRecord(**p["orf"])
                protein = ProteinRecord(
                    orf=orf,
                    aa_sequence=p["aa_sequence"],
                    aa_length=p["aa_length"],
                )
                proteins.append(protein)
        else:
            raise ValueError("Invalid protein JSON format")

        typer.echo(f"  Loaded {len(proteins)} protein(s)")

        typer.echo(f"\nInitializing ProstT5 encoder (model: {model})...")
        encoder = ProstT5ThreeDiEncoder(model_name=model, device=device)
        
        if multi_gpu:
            typer.echo(f"  Multi-GPU encoding: enabled")
            if parsed_gpu_ids:
                typer.echo(f"  GPU IDs: {parsed_gpu_ids}")
            else:
                typer.echo(f"  GPU IDs: auto-discover")
        else:
            typer.echo(f"  Using device: {encoder.device}")

        typer.echo(f"\nEncoding to 3Di tokens...")
        three_dis = encoder.encode_proteins(
            proteins,
            encoding_size,
            use_multi_gpu=multi_gpu,
            gpu_ids=parsed_gpu_ids,
        )
        typer.echo(f"  Encoded {len(three_dis)} sequence(s)")

        typer.echo(f"\nWriting results to: {output}")
        write_json(three_dis, output)

        typer.echo("✓ 3Di encoding complete!")

    except Exception as e:
        typer.echo(f"Error: {e}", err=True)
        traceback.print_exc()
        raise typer.Exit(3)
