import json
import os
import time
from typing import Any, List, Mapping, Optional, Union

import flux
import flux.job

import flux_mcp.utils as utils


def get_handle(uri: Optional[str] = None) -> flux.Flux:
    """Helper to get a Flux handle, optionally connecting to a remote URI."""
    if uri:
        return flux.Flux(uri)
    return flux.Flux()


def flux_submit_job(
    command: List[str],
    uri: Optional[str] = None,
    submit_async: bool = True,
    num_tasks: int = 1,
    cores_per_task: int = 1,
    gpus_per_task: Optional[int] = None,
    num_nodes: Optional[int] = None,
    exclusive: bool = False,
    duration: Optional[Union[int, float, str]] = None,
    environment: Optional[Mapping[str, str]] = None,
    env_expand: Optional[Mapping[str, str]] = None,
    cwd: Optional[str] = None,
    rlimits: Optional[Mapping[str, int]] = None,
    name: Optional[str] = None,
    input: Optional[Union[str, os.PathLike]] = None,
    output: Optional[Union[str, os.PathLike]] = None,
    error: Optional[Union[str, os.PathLike]] = None,
    label_io: bool = False,
    unbuffered: bool = False,
    queue: Optional[str] = None,
    bank: Optional[str] = None,
) -> str:
    """
    Creates a Jobspec from a command and submits it to Flux.

    Args:
        command: Command to execute (iterable of strings).
        uri: Optional Flux URI. If not provided, uses local instance.
        submit_async: Whether to submit the job asynchronously.
        num_tasks: Number of tasks to create.
        cores_per_task: Number of cores to allocate per task.
        gpus_per_task: Number of GPUs to allocate per task.
        num_nodes: Distribute allocated tasks across N individual nodes.
        exclusive: Always allocate nodes exclusively.
        duration: Time limit in Flux Standard Duration (str), seconds (int/float), or timedelta.
        environment: Mapping of environment variables for the job.
        env_expand: Mapping of environment variables containing mustache templates.
        cwd: Set the current working directory for the job.
        rlimits: Mapping of process resource limits (e.g. {"nofile": 12000}).
        name: Set a custom job name.
        input: Path to a file for job input.
        output: Path to a file for job output (stdout).
        error: Path to a file for job error (stderr).
        label_io: Label output with the source task IDs.
        unbuffered: Disable output buffering.
        queue: Set the queue for the job.
        bank: Set the bank for the job.

    Returns:
        JSON string containing the success status and Job ID or error message.
    """
    try:
        # Generate the Jobspec from the provided arguments
        jobspec = flux.job.JobspecV1.from_command(
            command=command,
            num_tasks=num_tasks,
            cores_per_task=cores_per_task,
            gpus_per_task=gpus_per_task,
            num_nodes=num_nodes,
            exclusive=exclusive,
            duration=duration,
            environment=environment,
            env_expand=env_expand,
            cwd=cwd,
            rlimits=rlimits,
            name=name,
            input=input,
            output=output,
            error=error,
            label_io=label_io,
            unbuffered=unbuffered,
            queue=queue,
            bank=bank,
        )
        h = get_handle(uri)

        # Submit the job
        if submit_async:
            future = flux.job.submit_async(h, jobspec)
            jobid = future.get_id()
        else:
            jobid = flux.job.submit(h, jobspec)
        return json.dumps({"success": True, "job_id": int(jobid), "uri": uri or "local"})

    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


def flux_submit_jobspec(jobspec: str, uri: Optional[str] = None, submit_async: bool = True) -> str:
    """
    Submits a job to Flux. An example jobspec requires attributes, tasks, resources, version
    and at least one referenced slot.
    {
        "version": 1,
        "resources": [
            {
                "type": "node",
                "count": 1,
                "with": [
                    {
                        "type": "slot",
                        "count": 1,
                        "label": "task",
                        "with": [{"type": "core", "count": 1}],
                    }
                ],
            }
        ],
        "tasks": [{"command": ["/bin/true"], "slot": "task", "count": {"per_slot": 1}}],
        "attributes": {"system": {"duration": 0}},
    }

    Args:
        jobspec: A valid JSON string or YAML string of the jobspec.
        uri: Optional Flux URI. If not provided, uses local instance.

    Returns:
        JSON string containing the new Job ID or error message.
    """
    try:
        # Ensure we got a json string
        jobspec = json.dumps(utils.load_jobspec(jobspec))
        h = get_handle(uri)

        # Submit the job
        if submit_async:
            future = flux.job.submit_async(h, jobspec)
            jobid = future.get_id()
        else:
            jobid = flux.job.submit(h, jobspec)

        # Return success with the integer ID
        return json.dumps({"success": True, "job_id": int(jobid), "uri": uri or "local"})

    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


def flux_cancel_job(job_id: Union[int, str], uri: Optional[str] = None) -> str:
    """
    Cancels a specific Flux job.

    Args:
        job_id: The ID of the job to cancel.
        uri: Optional Flux URI.
    """
    try:
        h = get_handle(uri)

        # Convert to proper integer ID
        job_id = flux.job.JobID(job_id)
        flux.job.cancel(h, job_id)

        return json.dumps(
            {"success": True, "message": f"Job {job_id} cancellation requested.", "job_id": job_id}
        )

    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


def flux_get_job_info(job_id: Union[int, str], uri: Optional[str] = None) -> str:
    """
    Retrieves status and information about a specific job.

    Args:
        job_id: The ID of the job.
        uri: Optional Flux URI.
    """
    try:
        h = get_handle(uri)
        id_int = flux.job.JobID(job_id)
        info = flux.job.get_job(h, id_int)
        return json.dumps(info)

    except EnvironmentError:
        return json.dumps({"success": False, "error": f"Job {job_id} not found."})
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


def flux_get_job_logs(job_id: Union[int, str], uri: Optional[str] = None, delay: int = 0) -> list:
    """
    Retrieves status and information about a specific job.

    Args:
        job_id: The ID of the job.
        uri: Optional Flux URI.
        delay: How long to wait (defaults to 0)
    """
    lines = []
    start = time.time()
    try:
        h = get_handle(uri)
        job_id = flux.job.JobID(job_id)
        for line in flux.job.event_watch(h, job_id, "guest.output"):
            if "data" in line.context:
                lines.append(line.context["data"])
            now = time.time()
            if delay is not None and (now - start) > delay:
                return lines
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})
    return lines
