"""Core job submission logic for clserve."""

import os
import random
import string
import tempfile
import subprocess
import logging
from dataclasses import dataclass

from jinja2 import Template

try:
    from importlib.resources import files
except ImportError:
    from importlib_resources import files

from clserve.configs import load_model_config, get_model_path
from clserve.config import get_account

logger = logging.getLogger(__name__)


@dataclass
class SubmitArgs:
    """Arguments for submitting a serving job."""

    model: str
    workers: int = 1
    nodes_per_worker: int = 1
    partition: str = "normal"
    environment: str = "sglang_gb200"
    tp_size: int = 1
    dp_size: int = 1
    ep_size: int = 1
    num_gpus_per_worker: int = 4
    cuda_graph_max_bs: int = 256
    grammar_backend: str = "llguidance"
    use_router: bool = False
    router_policy: str = "cache_aware"
    router_environment: str = "sglang_router"
    reasoning_parser: str = ""
    tool_call_parser: str = ""
    time_limit: str = "04:00:00"


def nanoid(length: int = 6) -> str:
    """Generate a random alphanumeric ID."""
    return "".join(random.choices(string.ascii_lowercase + string.digits, k=length))


def generate_job_name(model_path: str) -> str:
    """Generate a job name with clserve prefix and nanoid.

    Format: clserve_{nanoid}

    Args:
        model_path: Full model path (unused, kept for API compatibility)

    Returns:
        Job name with clserve prefix for easy filtering
    """
    return f"clserve_{nanoid()}"


def validate_args(args: SubmitArgs) -> None:
    """Validate submission arguments.

    Args:
        args: Submission arguments to validate

    Raises:
        ValueError: If arguments are invalid
    """
    if args.num_gpus_per_worker not in [1, 2, 4]:
        raise ValueError("--num-gpus-per-worker must be 1, 2, or 4")

    if args.num_gpus_per_worker < 4:
        if args.nodes_per_worker > 1:
            raise ValueError(
                "--num-gpus-per-worker < 4 requires --nodes-per-worker=1 "
                "(multi-node distributed not supported with split GPUs)"
            )
        if args.tp_size > args.num_gpus_per_worker:
            raise ValueError(
                f"--tp-size ({args.tp_size}) cannot be greater than "
                f"--num-gpus-per-worker ({args.num_gpus_per_worker})"
            )


def merge_with_config(args: SubmitArgs) -> SubmitArgs:
    """Merge user args with predefined model config if available.

    User-provided arguments take precedence over config defaults.

    Args:
        args: User-provided arguments

    Returns:
        SubmitArgs with config defaults applied for unset values
    """
    config = load_model_config(args.model)
    if config is None:
        # No predefined config, just resolve the model path
        model_path = get_model_path(args.model)
        return SubmitArgs(
            model=model_path,
            workers=args.workers,
            nodes_per_worker=args.nodes_per_worker,
            partition=args.partition,
            environment=args.environment,
            tp_size=args.tp_size,
            dp_size=args.dp_size,
            ep_size=args.ep_size,
            num_gpus_per_worker=args.num_gpus_per_worker,
            cuda_graph_max_bs=args.cuda_graph_max_bs,
            grammar_backend=args.grammar_backend,
            use_router=args.use_router,
            router_policy=args.router_policy,
            router_environment=args.router_environment,
            reasoning_parser=args.reasoning_parser,
            tool_call_parser=args.tool_call_parser,
            time_limit=args.time_limit,
        )

    # Use config values as defaults, but user args override
    return SubmitArgs(
        model=config.model_path,
        workers=args.workers if args.workers != 1 else config.workers,
        nodes_per_worker=(
            args.nodes_per_worker
            if args.nodes_per_worker != 1
            else config.nodes_per_worker
        ),
        partition=args.partition,
        environment=args.environment,
        tp_size=args.tp_size if args.tp_size != 1 else config.tp_size,
        dp_size=args.dp_size if args.dp_size != 1 else config.dp_size,
        ep_size=args.ep_size if args.ep_size != 1 else config.ep_size,
        num_gpus_per_worker=(
            args.num_gpus_per_worker
            if args.num_gpus_per_worker != 4
            else config.num_gpus_per_worker
        ),
        cuda_graph_max_bs=args.cuda_graph_max_bs,
        grammar_backend=args.grammar_backend,
        use_router=args.use_router or config.use_router,
        router_policy=args.router_policy
        if args.router_policy != "cache_aware"
        else config.router_policy,
        router_environment=args.router_environment,
        reasoning_parser=args.reasoning_parser or config.reasoning_parser,
        tool_call_parser=args.tool_call_parser or config.tool_call_parser,
        time_limit=args.time_limit,
    )


def render_job_script(args: SubmitArgs) -> str:
    """Render the SLURM job script from template.

    Args:
        args: Submission arguments

    Returns:
        Rendered job script content
    """
    template_content = (files("clserve") / "templates" / "job.jinja").read_text()
    template = Template(template_content)

    # Calculate total nodes
    total_nodes = args.nodes_per_worker * args.workers

    # Generate job name
    job_name = generate_job_name(args.model)

    # Format reasoning parser argument
    reasoning_parser_arg = ""
    if args.reasoning_parser:
        reasoning_parser_arg = f"--reasoning-parser {args.reasoning_parser}"

    # Format tool call parser argument
    tool_call_parser_arg = ""
    if args.tool_call_parser:
        tool_call_parser_arg = f"--tool-call-parser {args.tool_call_parser}"

    # Get home directory for log path
    home = os.path.expanduser("~")

    # Get cluster account from config or system
    cluster_account = get_account()

    return template.render(
        job_name=job_name,
        cluster_account=cluster_account,
        nodes=total_nodes,
        nodes_per_worker=args.nodes_per_worker,
        workers=args.workers,
        partition=args.partition,
        environment=args.environment,
        model_path=args.model,
        tp_size=args.tp_size,
        dp_size=args.dp_size,
        ep_size=args.ep_size,
        num_gpus_per_worker=args.num_gpus_per_worker,
        cuda_graph_max_bs=args.cuda_graph_max_bs,
        grammar_backend=args.grammar_backend,
        use_router=args.use_router,
        router_policy=args.router_policy,
        router_environment=args.router_environment,
        reasoning_parser=reasoning_parser_arg,
        tool_call_parser=tool_call_parser_arg,
        time_limit=args.time_limit,
        home=home,
    )


def submit_job(script_content: str) -> str:
    """Submit job script to SLURM.

    Args:
        script_content: Rendered job script

    Returns:
        Job ID

    Raises:
        RuntimeError: If job submission fails
    """
    with tempfile.NamedTemporaryFile(mode="w", suffix=".sh", delete=False) as f:
        f.write(script_content)
        script_path = f.name

    try:
        result = subprocess.run(
            ["sbatch", script_path],
            capture_output=True,
            text=True,
            check=True,
        )
        output_lines = result.stdout.strip().split("\n")
        job_id = output_lines[-1].split()[-1]
        return job_id
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"Failed to submit job: {e.stderr}")
    finally:
        os.unlink(script_path)


def ensure_clserve_dir() -> None:
    """Ensure the ~/.clserve/logs directory exists."""
    clserve_dir = os.path.expanduser("~/.clserve/logs")
    os.makedirs(clserve_dir, exist_ok=True)


def serve(args: SubmitArgs) -> str:
    """Submit a model serving job.

    This is the main entry point for submitting serving jobs.

    Args:
        args: Submission arguments

    Returns:
        Job ID of submitted job
    """
    # Ensure clserve directory exists
    ensure_clserve_dir()

    # Merge with predefined config if available
    merged_args = merge_with_config(args)

    # Validate arguments
    validate_args(merged_args)

    # Check if router should be auto-enabled
    processes_per_node = 4 // merged_args.num_gpus_per_worker
    total_processes = (
        merged_args.workers * merged_args.nodes_per_worker * processes_per_node
    )
    if total_processes > 1 and not merged_args.use_router:
        logger.warning(
            "Multiple worker processes detected but router is disabled. "
            "Consider using --use-router for load balancing."
        )

    # Render and submit job
    script = render_job_script(merged_args)
    job_id = submit_job(script)

    return job_id
