import secrets
from pathlib import Path

import typer
import yaml
from rich.console import Console
from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn

from src.aws import s3 as s3_ops
from src.commands import ai_hub, infra, train
from src.commands.utils import CONFIG_OPT, load_cfg
from src.config import GENERATED_STACK_PREFIX, Config, InfraConfig, S3Config

app = typer.Typer(
    help="qc-cli: End-to-end model managment for Qualcomm AI Hub.",
    no_args_is_help=True,
)
app.add_typer(infra.app, name="infra")
app.add_typer(train.app, name="train")
app.add_typer(ai_hub.app, name="ai-hub")

console = Console()


@app.command()
def init(
    output: str = typer.Option("config.yaml", help="Destination path for the config file"),
    force: bool = typer.Option(False, "--force", "-f", help="Overwrite an existing config file"),
) -> None:
    """Write a starter config.yaml to the current directory."""
    dest = Path(output)
    if dest.exists() and not force:
        console.print(f"[yellow]{dest} already exists.[/yellow] Use --force to overwrite.")
        raise typer.Exit(1)

    config = _new_isolated_config()
    dest.parent.mkdir(parents=True, exist_ok=True)
    config_data = config.model_dump(mode="json")
    config_data["sagemaker"].pop("role_name", None)
    with open(dest, "w") as f:
        yaml.safe_dump(config_data, f, sort_keys=False)

    console.print(f"[green]✓[/green] Config written to [bold]{dest}[/bold]")
    console.print(
        "Edit [cyan]sagemaker.training.image_uri[/cyan] before running training commands."
    )


def _new_isolated_config() -> Config:
    suffix = secrets.token_hex(6)
    namespace = f"{GENERATED_STACK_PREFIX}{suffix}"
    config = Config(infra=InfraConfig(stack_name=namespace))
    config.s3 = S3Config(bucket=f"{namespace}-data")
    return config


@app.command()
def upload(
    path: Path = typer.Argument(..., help="Local file or directory to upload"),
    s3_key: str | None = typer.Option(None, "--s3-key", help="S3 key for file uploads"),
    config: str = CONFIG_OPT,
) -> None:
    """Upload a local file or directory to S3."""
    cfg = load_cfg(config)

    if path.is_file():
        key = s3_key or f"{cfg.s3.data_prefix.rstrip('/')}/{path.name}"
        try:
            with console.status(f"Uploading {path.name}..."):
                uri = s3_ops.upload_file(cfg.aws.region, cfg.aws.profile, cfg.s3.bucket, str(path), key)
        except Exception as e:
            console.print(f"[red]Upload failed: {e}[/red]")
            raise typer.Exit(1)

        console.print(f"[green]✓[/green] {path.name} -> {uri}")
        return

    if path.is_dir():
        if s3_key is not None:
            console.print("[red]--s3-key can only be used when uploading a single file.[/red]")
            raise typer.Exit(1)

        files = [file for file in path.rglob("*") if file.is_file()]
        if not files:
            console.print("[yellow]No files found in directory.[/yellow]")
            raise typer.Exit(0)

        prefix = cfg.s3.data_prefix
        console.print(f"Uploading {len(files)} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/")
        try:
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                TaskProgressColumn(),
                console=console,
            ) as progress:
                task = progress.add_task("Uploading...", total=len(files))
                count = s3_ops.upload_dir(
                    cfg.aws.region,
                    cfg.aws.profile,
                    cfg.s3.bucket,
                    str(path),
                    prefix,
                    on_progress=lambda: progress.advance(task),
                )
        except Exception as e:
            console.print(f"[red]Upload failed: {e}[/red]")
            raise typer.Exit(1)

        console.print(f"[green]✓[/green] Uploaded {count} files to s3://{cfg.s3.bucket}/{prefix.rstrip('/')}/")
        return

    console.print(f"[red]Path not found: {path}[/red]")
    raise typer.Exit(1)