"""Main CLI interface for srunx."""

import argparse
import os
import sys
from pathlib import Path

from srunx.client import Slurm
from srunx.logging import (
    configure_cli_logging,
    configure_workflow_logging,
    get_logger,
)
from srunx.models import Job, JobEnvironment, JobResource, ShellJob, Workflow
from srunx.workflows.runner import WorkflowRunner

logger = get_logger(__name__)


def create_job_parser() -> argparse.ArgumentParser:
    """Create argument parser for job submission."""
    parser = argparse.ArgumentParser(
        description="Submit SLURM jobs with various configurations",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    # Required arguments
    parser.add_argument(
        "command",
        nargs="+",
        help="Command to execute in the SLURM job",
    )

    # Job configuration
    parser.add_argument(
        "--name",
        "--job-name",
        type=str,
        default="job",
        help="Job name (default: %(default)s)",
    )
    parser.add_argument(
        "--log-dir",
        type=str,
        default=os.getenv("SLURM_LOG_DIR", "logs"),
        help="Log directory (default: %(default)s)",
    )
    parser.add_argument(
        "--work-dir",
        "--chdir",
        type=str,
        help="Working directory for the job",
    )

    # Resource configuration
    resource_group = parser.add_argument_group("Resource Options")
    resource_group.add_argument(
        "-N",
        "--nodes",
        type=int,
        default=1,
        help="Number of nodes (default: %(default)s)",
    )
    resource_group.add_argument(
        "--gpus-per-node",
        type=int,
        default=0,
        help="Number of GPUs per node (default: %(default)s)",
    )
    resource_group.add_argument(
        "--ntasks-per-node",
        type=int,
        default=1,
        help="Number of tasks per node (default: %(default)s)",
    )
    resource_group.add_argument(
        "--cpus-per-task",
        type=int,
        default=1,
        help="Number of CPUs per task (default: %(default)s)",
    )
    resource_group.add_argument(
        "--memory",
        "--mem",
        type=str,
        help="Memory per node (e.g., '32GB', '1TB')",
    )
    resource_group.add_argument(
        "--time",
        "--time-limit",
        type=str,
        help="Time limit (e.g., '1:00:00', '30:00', '1-12:00:00')",
    )

    # Environment configuration
    env_group = parser.add_argument_group("Environment Options")
    env_group.add_argument(
        "--conda",
        type=str,
        help="Conda environment name",
    )
    env_group.add_argument(
        "--venv",
        type=str,
        help="Virtual environment path",
    )
    env_group.add_argument(
        "--sqsh",
        type=str,
        help="SquashFS image path",
    )
    env_group.add_argument(
        "--env",
        action="append",
        dest="env_vars",
        help="Environment variable KEY=VALUE (can be used multiple times)",
    )

    # Execution options
    exec_group = parser.add_argument_group("Execution Options")
    exec_group.add_argument(
        "--template",
        type=str,
        help="Path to custom SLURM template file",
    )
    exec_group.add_argument(
        "--wait",
        action="store_true",
        help="Wait for job completion",
    )
    exec_group.add_argument(
        "--poll-interval",
        type=int,
        default=30,
        help="Polling interval in seconds when waiting (default: %(default)s)",
    )

    # Logging options
    log_group = parser.add_argument_group("Logging Options")
    log_group.add_argument(
        "--log-level",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        default="INFO",
        help="Set logging level (default: %(default)s)",
    )
    log_group.add_argument(
        "--quiet",
        "-q",
        action="store_true",
        help="Only show warnings and errors",
    )

    return parser


def create_status_parser() -> argparse.ArgumentParser:
    """Create argument parser for job status."""
    parser = argparse.ArgumentParser(
        description="Check SLURM job status",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument(
        "job_id",
        type=int,
        help="SLURM job ID to check",
    )

    return parser


def create_list_parser() -> argparse.ArgumentParser:
    """Create argument parser for listing jobs."""
    parser = argparse.ArgumentParser(
        description="List SLURM jobs",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument(
        "--user",
        "-u",
        type=str,
        help="List jobs for specific user (default: current user)",
    )

    return parser


def create_cancel_parser() -> argparse.ArgumentParser:
    """Create argument parser for job cancellation."""
    parser = argparse.ArgumentParser(
        description="Cancel SLURM job",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument(
        "job_id",
        type=int,
        help="SLURM job ID to cancel",
    )

    return parser


def create_main_parser() -> argparse.ArgumentParser:
    """Create main argument parser with subcommands."""
    parser = argparse.ArgumentParser(
        description="srunx - Python library for SLURM job management",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    # Global options
    parser.add_argument(
        "--log-level",
        "-l",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        default="INFO",
        help="Set logging level (default: %(default)s)",
    )
    parser.add_argument(
        "--quiet",
        "-q",
        action="store_true",
        help="Only show warnings and errors",
    )

    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # Submit command (default)
    submit_parser = subparsers.add_parser("submit", help="Submit a SLURM job")
    submit_parser.set_defaults(func=cmd_submit)
    _copy_parser_args(create_job_parser(), submit_parser)

    # Status command
    status_parser = subparsers.add_parser("status", help="Check job status")
    status_parser.set_defaults(func=cmd_status)
    _copy_parser_args(create_status_parser(), status_parser)

    # List command
    list_parser = subparsers.add_parser("list", help="List jobs")
    list_parser.set_defaults(func=cmd_list)
    _copy_parser_args(create_list_parser(), list_parser)

    # Cancel command
    cancel_parser = subparsers.add_parser("cancel", help="Cancel job")
    cancel_parser.set_defaults(func=cmd_cancel)
    _copy_parser_args(create_cancel_parser(), cancel_parser)

    # Flow command
    flow_parser = subparsers.add_parser("flow", help="Workflow management")
    flow_parser.set_defaults(func=None)  # Will be overridden by subcommands

    # Flow subcommands
    flow_subparsers = flow_parser.add_subparsers(
        dest="flow_command", help="Flow commands"
    )

    # Flow run command
    flow_run_parser = flow_subparsers.add_parser("run", help="Execute workflow")
    flow_run_parser.set_defaults(func=cmd_flow_run)
    flow_run_parser.add_argument(
        "yaml_file",
        type=str,
        help="Path to YAML workflow definition file",
    )
    flow_run_parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be executed without running jobs",
    )

    # Flow validate command
    flow_validate_parser = flow_subparsers.add_parser(
        "validate", help="Validate workflow"
    )
    flow_validate_parser.set_defaults(func=cmd_flow_validate)
    flow_validate_parser.add_argument(
        "yaml_file",
        type=str,
        help="Path to YAML workflow definition file",
    )

    return parser


def _copy_parser_args(
    source_parser: argparse.ArgumentParser, target_parser: argparse.ArgumentParser
) -> None:
    """Copy arguments from source parser to target parser."""
    for action in source_parser._actions:
        if action.dest == "help":
            continue
        target_parser._add_action(action)


def _parse_env_vars(env_var_list: list[str] | None) -> dict[str, str]:
    """Parse environment variables from list of KEY=VALUE strings."""
    env_vars = {}
    if env_var_list:
        for env_var in env_var_list:
            if "=" in env_var:
                key, value = env_var.split("=", 1)
                env_vars[key] = value
            else:
                logger.warning(f"Invalid environment variable format: {env_var}")
    return env_vars


def cmd_submit(args: argparse.Namespace) -> None:
    """Handle job submission command."""
    try:
        # Parse environment variables
        env_vars = _parse_env_vars(getattr(args, "env_vars", None))

        # Create job configuration
        resources = JobResource(
            nodes=args.nodes,
            gpus_per_node=args.gpus_per_node,
            ntasks_per_node=args.ntasks_per_node,
            cpus_per_task=args.cpus_per_task,
            memory_per_node=getattr(args, "memory", None),
            time_limit=getattr(args, "time", None),
        )

        environment = JobEnvironment(
            conda=getattr(args, "conda", None),
            venv=getattr(args, "venv", None),
            sqsh=getattr(args, "sqsh", None),
            env_vars=env_vars,
        )

        job_data = {
            "name": args.name,
            "command": args.command,
            "resources": resources,
            "environment": environment,
            "log_dir": args.log_dir,
        }

        if args.work_dir is not None:
            job_data["work_dir"] = args.work_dir

        job = Job.model_validate(job_data)

        # Submit job
        client = Slurm()
        submitted_job = client.run(job, getattr(args, "template", None))

        logger.info(f"Submitted job {submitted_job.job_id}: {submitted_job.name}")

        # Wait for completion if requested
        if getattr(args, "wait", False):
            logger.info(f"Waiting for job {submitted_job.job_id} to complete...")
            completed_job = client.monitor(
                submitted_job, poll_interval=args.poll_interval
            )
            status_str = (
                completed_job.status.value if completed_job.status else "Unknown"
            )
            logger.info(
                f"Job {submitted_job.job_id} completed with status: {status_str}"
            )

    except Exception as e:
        logger.error(f"Error submitting job: {e}")
        sys.exit(1)


def cmd_status(args: argparse.Namespace) -> None:
    """Handle job status command."""
    try:
        client = Slurm()
        job = client.retrieve(args.job_id)

        logger.info(f"Job ID: {job.job_id}")
        logger.info(f"Name: {job.name}")
        if job.status:
            logger.info(f"Status: {job.status.value}")
        else:
            logger.info("Status: Unknown")

    except Exception as e:
        logger.error(f"Error getting job status: {e}")
        sys.exit(1)


def cmd_list(args: argparse.Namespace) -> None:
    """Handle job listing command."""
    try:
        client = Slurm()
        jobs = client.list(getattr(args, "user", None))

        if not jobs:
            logger.info("No jobs found")
            return

        logger.info(f"{'Job ID':<10} {'Name':<20} {'Status':<12}")
        logger.info("-" * 45)
        for job in jobs:
            status_str = job.status.value if job.status else "Unknown"
            logger.info(f"{job.job_id:<10} {job.name:<20} {status_str:<12}")

    except Exception as e:
        logger.error(f"Error listing jobs: {e}")
        sys.exit(1)


def cmd_cancel(args: argparse.Namespace) -> None:
    """Handle job cancellation command."""
    try:
        client = Slurm()
        client.cancel(args.job_id)
        logger.info(f"Cancelled job {args.job_id}")

    except Exception as e:
        logger.error(f"Error cancelling job: {e}")
        sys.exit(1)


def cmd_flow_run(args: argparse.Namespace) -> None:
    """Handle flow run command."""
    # Configure logging for workflow execution
    configure_workflow_logging(level=getattr(args, "log_level", "INFO"))

    try:
        yaml_file = Path(args.yaml_file)
        if not yaml_file.exists():
            logger.error(f"Workflow file not found: {args.yaml_file}")
            sys.exit(1)

        runner = WorkflowRunner()

        # Load workflow for validation
        workflow = runner.load_from_yaml(yaml_file)
        logger.info(
            f"Loaded workflow '{workflow.name}' with {len(workflow.tasks)} tasks"
        )

        # Validate dependencies
        _validate_workflow_dependencies(workflow)

        if args.dry_run:
            _show_workflow_plan(workflow)
            return

        # Execute workflow
        logger.info("Starting workflow execution")
        results = runner.execute_workflow(workflow)

        logger.info("Workflow execution completed successfully")
        logger.info("Job Results:")
        for task_name, job in results.items():
            if hasattr(job, "job_id") and job.job_id:
                logger.info(f"  {task_name}: Job ID {job.job_id}")
            else:
                logger.info(f"  {task_name}: {job}")

    except Exception as e:
        logger.error(f"Workflow execution failed: {e}")
        sys.exit(1)


def cmd_flow_validate(args: argparse.Namespace) -> None:
    """Handle flow validate command."""
    # Configure logging for workflow validation
    configure_workflow_logging(level=getattr(args, "log_level", "INFO"))

    try:
        yaml_file = Path(args.yaml_file)
        if not yaml_file.exists():
            logger.error(f"Workflow file not found: {args.yaml_file}")
            sys.exit(1)

        runner = WorkflowRunner()

        # Load workflow for validation
        workflow = runner.load_from_yaml(yaml_file)
        logger.info(
            f"Loaded workflow '{workflow.name}' with {len(workflow.tasks)} tasks"
        )

        # Validate dependencies
        _validate_workflow_dependencies(workflow)

        logger.info("Workflow validation successful")

    except Exception as e:
        logger.error(f"Workflow validation failed: {e}")
        sys.exit(1)


def _validate_workflow_dependencies(workflow: Workflow) -> None:
    """Validate workflow task dependencies."""
    task_names = {task.name for task in workflow.tasks}

    for task in workflow.tasks:
        for dependency in task.depends_on:
            if dependency not in task_names:
                raise ValueError(
                    f"Task '{task.name}' depends on unknown task '{dependency}'"
                )

    # Check for circular dependencies (simple check)
    visited = set()
    rec_stack = set()

    def has_cycle(task_name: str) -> bool:
        if task_name in rec_stack:
            return True
        if task_name in visited:
            return False

        visited.add(task_name)
        rec_stack.add(task_name)

        task = workflow.get_task(task_name)
        if task:
            for dependency in task.depends_on:
                if has_cycle(dependency):
                    return True

        rec_stack.remove(task_name)
        return False

    for task in workflow.tasks:
        if has_cycle(task.name):
            raise ValueError(
                f"Circular dependency detected involving task '{task.name}'"
            )


def _show_workflow_plan(workflow: Workflow) -> None:
    """Show workflow execution plan."""
    msg = """\
Workflow execution plan:
    Workflow: {workflow.name}
    Tasks: {len(workflow.tasks)}
"""

    for task in workflow.tasks:
        msg += f"\t\tTask: {task.name}\n"
        if isinstance(task.job, Job):
            msg += f"\t\t\tCommand: {' '.join(task.job.command or [])}\n"
            msg += f"\t\t\tResources: {task.job.resources.nodes} nodes, {task.job.resources.gpus_per_node} GPUs/node\n"
            if task.job.environment.conda:
                msg += f"\t\t\tConda env: {task.job.environment.conda}\n"
            if task.job.environment.sqsh:
                msg += f"\t\t\tSqsh: {task.job.environment.sqsh}\n"
            if task.job.environment.venv:
                msg += f"\t\t\tVenv: {task.job.environment.venv}\n"
        elif isinstance(task.job, ShellJob):
            msg += f"\t\t\tPath: {task.job.path}\n"
        if task.depends_on:
            msg += f"\t\t\tDependencies: {', '.join(task.depends_on)}\n"
        if task.async_execution:
            msg += "\t\t\tExecution: asynchronous\n"

    logger.info(msg)


def main() -> None:
    """Main entry point for the CLI."""
    parser = create_main_parser()
    args = parser.parse_args()

    # Configure logging
    log_level = getattr(args, "log_level", "INFO")
    quiet = getattr(args, "quiet", False)
    configure_cli_logging(level=log_level, quiet=quiet)

    # If no command specified, default to submit behavior for backward compatibility
    if not hasattr(args, "func") or args.func is None:
        # Check if this is a flow command without subcommand
        if hasattr(args, "command") and args.command == "flow":
            if not hasattr(args, "flow_command") or args.flow_command is None:
                logger.error("Flow command requires a subcommand (run or validate)")
                parser.print_help()
                sys.exit(1)
        else:
            # Try to parse as submit command
            submit_parser = create_job_parser()
            try:
                submit_args = submit_parser.parse_args()
                cmd_submit(submit_args)
            except SystemExit:
                parser.print_help()
                sys.exit(1)
    else:
        args.func(args)


if __name__ == "__main__":
    main()
