import builtins
from enum import Enum
from typing import List, Union

from picle.models import Outputters, PipeFunctionsModel
from pydantic import (
    BaseModel,
    Field,
    StrictBool,
    StrictInt,
    StrictStr,
)

from ..common import ClientRunJobArgs, log_error_or_result


class NornirTaskEnum(str, Enum):
    cli = "cli"
    parse = "parse"
    cfg = "cfg"
    task = "task"
    test = "test"
    network = "network"
    diagram = "diagram"


class ListJobsModel(ClientRunJobArgs):
    workers: StrictStr = Field("all", description="Workers to return jobs for")
    last: StrictInt = Field(
        None, description="Return last N completed and last N pending jobs"
    )
    pending: StrictBool = Field(
        True, description="Return pending jobs", json_schema_extra={"presence": True}
    )
    completed: StrictBool = Field(
        True, description="Return completed jobs", json_schema_extra={"presence": True}
    )
    client: StrictStr = Field(None, description="Client name to return jobs for")
    uuid: StrictStr = Field(None, description="Job UUID to return")
    task: NornirTaskEnum = Field(None, description="Task name to return jobs for")

    @staticmethod
    def source_client():
        return ["self"]

    @staticmethod
    def source_workers():
        NFCLIENT = builtins.NFCLIENT
        reply = NFCLIENT.mmi(
            "mmi.service.broker", "show_workers", kwargs={"service": "nornir"}
        )
        workers = [i["name"] for i in reply["results"]]

        return ["all", "any"] + workers

    @staticmethod
    def run(*args, **kwargs):
        NFCLIENT = builtins.NFCLIENT
        workers = kwargs.pop("workers", "all")
        timeout = kwargs.pop("timeout", 600)
        verbose_result = kwargs.pop("verbose_result", False)
        nowait = kwargs.pop("nowait", False)

        if kwargs.get("client") == "self":
            kwargs["client"] = NFCLIENT.zmq_name

        result = NFCLIENT.run_job(
            "nornir",
            "job_list",
            workers=workers,
            args=args,
            kwargs=kwargs,
            timeout=timeout,
            nowait=nowait,
        )

        if nowait:
            return result, Outputters.outputter_nested

        result = log_error_or_result(result, verbose_result=verbose_result)

        ret = []
        for worker_name, worker_results in result.items():
            ret.extend(worker_results)

        return ret

    class PicleConfig:
        outputter = Outputters.outputter_rich_table


class JobDetailsModel(ClientRunJobArgs):
    uuid: StrictStr = Field(..., description="Job UUID")
    workers: Union[StrictStr, List[StrictStr]] = Field(
        "all", description="Workers to return jobs for"
    )
    result: StrictBool = Field(
        True, description="Return job result", json_schema_extra={"presence": True}
    )
    events: StrictBool = Field(
        True, description="Return job events", json_schema_extra={"presence": True}
    )

    @staticmethod
    def source_workers():
        NFCLIENT = builtins.NFCLIENT
        reply = NFCLIENT.mmi(
            "mmi.service.broker", "show_workers", kwargs={"service": "nornir"}
        )
        workers = [i["name"] for i in reply["results"]]

        return ["all", "any"] + workers

    @staticmethod
    def run(*args, **kwargs):
        NFCLIENT = builtins.NFCLIENT
        workers = kwargs.pop("workers", "all")
        timeout = kwargs.pop("timeout", 600)
        verbose_result = kwargs.pop("verbose_result", False)
        nowait = kwargs.pop("nowait", False)

        result = NFCLIENT.run_job(
            "nornir",
            "job_details",
            workers=workers,
            args=args,
            kwargs=kwargs,
            timeout=timeout,
            nowait=nowait,
        )

        if nowait:
            return result, Outputters.outputter_nested

        result = log_error_or_result(result, verbose_result=verbose_result)

        return result

    class PicleConfig:
        outputter = Outputters.outputter_nested
        pipe = PipeFunctionsModel


class NornirJobsShell(BaseModel):
    summary: ListJobsModel = Field(None, description="List jobs")
    details: JobDetailsModel = Field(None, description="Show job details")

    class PicleConfig:
        pass
