"""Hanzo ML - Machine learning platform CLI.

End-to-end MLOps: notebooks, pipelines, training, serving, registry.
"""

import click
from rich import box
from rich.table import Table
from rich.panel import Panel

from ..utils.output import console


@click.group(name="ml")
def ml_group():
    """Hanzo ML - End-to-end machine learning platform.

    \b
    Develop:
      hanzo ml notebooks list      # List notebooks
      hanzo ml notebooks create    # Create notebook

    \b
    Train:
      hanzo ml training list       # List training jobs
      hanzo ml training create     # Start training job
      hanzo ml pipelines run       # Run ML pipeline

    \b
    Serve:
      hanzo ml serving list        # List deployments
      hanzo ml serving deploy      # Deploy model

    \b
    Manage:
      hanzo ml registry list       # List model versions
      hanzo ml registry push       # Push model to registry
    """
    pass


# ============================================================================
# Notebooks
# ============================================================================

@ml_group.group()
def notebooks():
    """Manage Jupyter notebooks."""
    pass


@notebooks.command(name="list")
def notebooks_list():
    """List all notebooks."""
    table = Table(title="Notebooks", box=box.ROUNDED)
    table.add_column("Name", style="cyan")
    table.add_column("Instance", style="white")
    table.add_column("Status", style="green")
    table.add_column("GPU", style="yellow")
    table.add_column("Created", style="dim")

    console.print(table)
    console.print("[dim]No notebooks found. Create one with 'hanzo ml notebooks create'[/dim]")


@notebooks.command(name="create")
@click.option("--name", "-n", prompt=True, help="Notebook name")
@click.option("--instance", "-i", default="cpu-small", help="Instance type")
@click.option("--gpu", is_flag=True, help="Enable GPU")
def notebooks_create(name: str, instance: str, gpu: bool):
    """Create a new notebook instance."""
    console.print(f"[cyan]Creating notebook '{name}'...[/cyan]")
    console.print(f"  Instance: {instance}")
    console.print(f"  GPU: {'Yes' if gpu else 'No'}")
    console.print()
    console.print(f"[green]✓[/green] Notebook '{name}' created")
    console.print(f"[dim]Access at: https://notebooks.hanzo.ai/{name}[/dim]")


@notebooks.command(name="start")
@click.argument("name")
def notebooks_start(name: str):
    """Start a notebook."""
    console.print(f"[green]✓[/green] Notebook '{name}' started")


@notebooks.command(name="stop")
@click.argument("name")
def notebooks_stop(name: str):
    """Stop a notebook."""
    console.print(f"[green]✓[/green] Notebook '{name}' stopped")


@notebooks.command(name="delete")
@click.argument("name")
def notebooks_delete(name: str):
    """Delete a notebook."""
    from rich.prompt import Confirm
    if not Confirm.ask(f"[red]Delete notebook '{name}'?[/red]"):
        return
    console.print(f"[green]✓[/green] Notebook '{name}' deleted")


# ============================================================================
# Training
# ============================================================================

@ml_group.group()
def training():
    """Manage training jobs."""
    pass


@training.command(name="list")
@click.option("--status", type=click.Choice(["running", "completed", "failed", "all"]), default="all")
def training_list(status: str):
    """List training jobs."""
    table = Table(title="Training Jobs", box=box.ROUNDED)
    table.add_column("ID", style="cyan")
    table.add_column("Name", style="white")
    table.add_column("Framework", style="white")
    table.add_column("Status", style="green")
    table.add_column("Duration", style="dim")
    table.add_column("GPU", style="yellow")

    console.print(table)


@training.command(name="create")
@click.option("--name", "-n", prompt=True, help="Job name")
@click.option("--framework", "-f", type=click.Choice(["pytorch", "tensorflow", "xgboost"]), default="pytorch")
@click.option("--script", "-s", required=True, help="Training script path")
@click.option("--gpu", "-g", default="1", help="Number of GPUs")
@click.option("--instance", "-i", default="gpu-a10g", help="Instance type")
def training_create(name: str, framework: str, script: str, gpu: str, instance: str):
    """Create a training job."""
    console.print(f"[cyan]Creating training job '{name}'...[/cyan]")
    console.print(f"  Framework: {framework}")
    console.print(f"  Script: {script}")
    console.print(f"  Instance: {instance} ({gpu} GPUs)")
    console.print()
    console.print(f"[green]✓[/green] Training job '{name}' started")


@training.command(name="logs")
@click.argument("job_id")
@click.option("--follow", "-f", is_flag=True, help="Follow logs")
def training_logs(job_id: str, follow: bool):
    """View training job logs."""
    console.print(f"[cyan]Logs for job {job_id}:[/cyan]")


@training.command(name="stop")
@click.argument("job_id")
def training_stop(job_id: str):
    """Stop a training job."""
    console.print(f"[green]✓[/green] Training job '{job_id}' stopped")


# ============================================================================
# Pipelines
# ============================================================================

@ml_group.group()
def pipelines():
    """Manage ML pipelines."""
    pass


@pipelines.command(name="list")
def pipelines_list():
    """List all pipelines."""
    table = Table(title="ML Pipelines", box=box.ROUNDED)
    table.add_column("Name", style="cyan")
    table.add_column("Version", style="white")
    table.add_column("Status", style="green")
    table.add_column("Last Run", style="dim")

    console.print(table)


@pipelines.command(name="create")
@click.option("--name", "-n", prompt=True, help="Pipeline name")
@click.option("--file", "-f", required=True, help="Pipeline definition file")
def pipelines_create(name: str, file: str):
    """Create a pipeline from definition file."""
    console.print(f"[green]✓[/green] Pipeline '{name}' created from {file}")


@pipelines.command(name="run")
@click.argument("pipeline_name")
@click.option("--params", "-p", help="JSON parameters")
def pipelines_run(pipeline_name: str, params: str):
    """Run a pipeline."""
    console.print(f"[cyan]Running pipeline '{pipeline_name}'...[/cyan]")
    console.print(f"[green]✓[/green] Pipeline run started")


# ============================================================================
# Serving
# ============================================================================

@ml_group.group()
def serving():
    """Manage model serving."""
    pass


@serving.command(name="list")
def serving_list():
    """List model deployments."""
    table = Table(title="Model Deployments", box=box.ROUNDED)
    table.add_column("Name", style="cyan")
    table.add_column("Model", style="white")
    table.add_column("Version", style="white")
    table.add_column("Status", style="green")
    table.add_column("Replicas", style="dim")
    table.add_column("Endpoint", style="dim")

    console.print(table)


@serving.command(name="deploy")
@click.option("--name", "-n", prompt=True, help="Deployment name")
@click.option("--model", "-m", required=True, help="Model path or registry URI")
@click.option("--replicas", "-r", default=1, help="Number of replicas")
@click.option("--gpu", is_flag=True, help="Enable GPU inference")
def serving_deploy(name: str, model: str, replicas: int, gpu: bool):
    """Deploy a model for inference."""
    console.print(f"[cyan]Deploying model '{name}'...[/cyan]")
    console.print(f"  Model: {model}")
    console.print(f"  Replicas: {replicas}")
    console.print(f"  GPU: {'Yes' if gpu else 'No'}")
    console.print()
    console.print(f"[green]✓[/green] Model deployed")
    console.print(f"[dim]Endpoint: https://serving.hanzo.ai/{name}[/dim]")


@serving.command(name="scale")
@click.argument("name")
@click.option("--replicas", "-r", required=True, type=int, help="Target replicas")
def serving_scale(name: str, replicas: int):
    """Scale a deployment."""
    console.print(f"[green]✓[/green] Deployment '{name}' scaled to {replicas} replicas")


@serving.command(name="delete")
@click.argument("name")
def serving_delete(name: str):
    """Delete a deployment."""
    console.print(f"[green]✓[/green] Deployment '{name}' deleted")


# ============================================================================
# Registry
# ============================================================================

@ml_group.group()
def registry():
    """Manage model registry."""
    pass


@registry.command(name="list")
@click.option("--model", "-m", help="Filter by model name")
def registry_list(model: str):
    """List registered models."""
    table = Table(title="Model Registry", box=box.ROUNDED)
    table.add_column("Model", style="cyan")
    table.add_column("Version", style="white")
    table.add_column("Stage", style="green")
    table.add_column("Framework", style="dim")
    table.add_column("Created", style="dim")

    console.print(table)


@registry.command(name="push")
@click.option("--name", "-n", required=True, help="Model name")
@click.option("--path", "-p", required=True, help="Model path")
@click.option("--framework", "-f", default="pytorch", help="Framework")
@click.option("--version", "-v", help="Version (auto-incremented if not specified)")
def registry_push(name: str, path: str, framework: str, version: str):
    """Push a model to the registry."""
    console.print(f"[cyan]Pushing model '{name}'...[/cyan]")
    console.print(f"  Path: {path}")
    console.print(f"  Framework: {framework}")
    console.print()
    console.print(f"[green]✓[/green] Model '{name}' pushed to registry")


@registry.command(name="pull")
@click.argument("model_uri")
@click.option("--output", "-o", default=".", help="Output directory")
def registry_pull(model_uri: str, output: str):
    """Pull a model from the registry."""
    console.print(f"[cyan]Pulling model {model_uri}...[/cyan]")
    console.print(f"[green]✓[/green] Model downloaded to {output}")


@registry.command(name="promote")
@click.argument("model_uri")
@click.option("--stage", "-s", type=click.Choice(["staging", "production"]), required=True)
def registry_promote(model_uri: str, stage: str):
    """Promote a model version to a stage."""
    console.print(f"[green]✓[/green] Model promoted to {stage}")
