"""Run command - Submit tasks from YAML configuration."""

import json
from pathlib import Path
from typing import Tuple

import click
import yaml

from flow import Flow, TaskConfig, ValidationError
from flow.errors import AuthenticationError

from .base import BaseCommand, console
from .utils import display_config, wait_for_task


class RunCommand(BaseCommand):
    """Submit a task from YAML configuration."""

    @property
    def name(self) -> str:
        return "run"

    @property
    def help(self) -> str:
        return "Submit a task from YAML configuration"

    def get_command(self) -> click.Command:
        @click.command(name=self.name, help=self.help)
        @click.argument("config_file", type=click.Path(exists=True))
        @click.option(
            "--wait/--no-wait",
            default=True,
            help="Wait for task to start running (default: wait)"
        )
        @click.option("--dry-run", is_flag=True, help="Validate configuration without submitting")
        @click.option("--watch", is_flag=True, help="Watch task progress interactively")
        @click.option("--json", "output_json", is_flag=True, help="Output JSON for automation")
        @click.option(
            "--mount",
            multiple=True,
            help="Mount storage (format: source or target=source). Auto-mounts: s3://→/data, volume://→/mnt",
        )
        def run(config_file: str, wait: bool, dry_run: bool, watch: bool, output_json: bool, mount: Tuple[str, ...]):
            """Submit a task from YAML configuration.

            CONFIG_FILE: Path to YAML configuration file

            \b
            Examples:
                flow run job.yaml              # Submit and wait
                flow run job.yaml --no-wait    # Submit and exit
                flow run job.yaml --watch      # Watch progress
                flow run job.yaml --json       # JSON output
                flow run job.yaml --mount s3://bucket/dataset  # Mount S3 bucket
                flow run job.yaml --mount /data=s3://bucket/dataset  # Mount to specific path
                flow run job.yaml --mount /checkpoints=volume://model-v1  # Mount volume
            """
            self._execute(config_file, wait, dry_run, watch, output_json, mount)

        return run

    def _execute(self, config_file: str, wait: bool, dry_run: bool,
                 watch: bool, output_json: bool, mount: Tuple[str, ...]) -> None:
        """Execute the run command."""
        try:
            # Load and display config
            config = TaskConfig.from_yaml(config_file)

            # Parse --mount flags
            mount_dict = None
            if mount:
                mount_dict = {}
                for mount_spec in mount:
                    if "=" in mount_spec:
                        # Format: target=source
                        target, source = mount_spec.split("=", 1)
                        mount_dict[target] = source
                    else:
                        # Format: source (auto-generate target)
                        source = mount_spec
                        if source.startswith("s3://"):
                            target = "/data"
                        elif source.startswith("volume://"):
                            target = "/mnt"
                        else:
                            # For local paths, use the basename
                            target = f"/mnt/{Path(source).name}"
                        mount_dict[target] = source

            if not output_json:
                display_config(config.model_dump())
                if mount_dict:
                    console.print("\n[bold]Mounts:[/bold]")
                    for target, source in mount_dict.items():
                        console.print(f"  {target} → {source}")

            if dry_run:
                if output_json:
                    result = {"status": "valid", "config": config.model_dump()}
                    if mount_dict:
                        result["mounts"] = mount_dict
                    console.print(json.dumps(result))
                else:
                    console.print("\n[green]✓[/green] Configuration is valid")
                return

            # Submit task
            flow_client = Flow()

            if not output_json:
                console.print("\n[cyan]Submitting task...[/cyan]")

            task = flow_client.run(config, mounts=mount_dict)

            if output_json:
                result = {"task_id": task.task_id, "status": "submitted"}
                if wait:
                    status = wait_for_task(flow_client, task.task_id, watch=False, json_output=True)
                    result["status"] = status.state
                    result["details"] = status.model_dump() if hasattr(status, 'model_dump') else status.dict()
                console.print(json.dumps(result))
                return

            if not wait:
                console.print(f"\nTask submitted with ID: [cyan]{task.task_id}[/cyan]")
                console.print(f"Check status with: [dim]flow status {task.task_id}[/dim]")
                return

            # Wait for task to start
            status = wait_for_task(flow_client, task.task_id, watch=watch, json_output=False)

            if status.state == "running":
                console.print("\n[green]✓[/green] Task is running!")
                if status.ssh_command:
                    console.print(f"\nConnect with: [cyan]{status.ssh_command}[/cyan]")
                console.print(f"Stream logs with: [dim]flow logs {task_id} --follow[/dim]")
            elif status.state == "failed":
                console.print("\n[red]✗[/red] Task failed to start")
                if status.error:
                    console.print(f"Error: {status.error}")
            elif status.state == "cancelled":
                console.print("\n[yellow]![/yellow] Task was cancelled")

        except AuthenticationError:
            self.handle_auth_error()
        except ValidationError as e:
            self.handle_error(f"Invalid configuration: {e}")
        except FileNotFoundError:
            self.handle_error(f"Configuration file not found: {config_file}")
        except yaml.YAMLError as e:
            self.handle_error(f"Invalid YAML: {e}")
        except Exception as e:
            self.handle_error(str(e))


# Export command instance
command = RunCommand()
