# job.py
#
# Copyright Uberware. All Rights Reserved

import json
import logging
import os
import typing
import uuid
from dataclasses import dataclass
from datetime import datetime
from functools import cached_property
from tempfile import NamedTemporaryFile
from uuid import UUID

from .component import (
    NameOrID,
    NameOrIDList,
    abort,
    datetime_to_rtime,
    ensure_list,
    extract_datetime,
    find_value,
    run_component,
    run_simple_command,
)
from .config import SMEDGE

"""Manage Jobs"""

SUBMIT = SMEDGE / "Submit"
"""Location of the Submit executable, uses config.SMEDGE by default"""

JOB = SMEDGE / "Job"
"""Location of the Job executable, uses config.SMEDGE by default"""

JobType = typing.Dict[str, any]
"""Type used for a Job"""

# Holds the product cache
_product_cache = {}

# Use Python standard logging
logger = logging.getLogger("smedge")


@dataclass(frozen=True)
class HistoryElement:
    """One Job History event"""

    name: str
    work_id: UUID
    time: datetime
    status: int
    engine_id: UUID
    note: str


@dataclass(frozen=True)
class Parameter:
    """One parameter in a product"""

    type: str
    flags: typing.List[str]
    settings: typing.Dict[str, str]
    choices: typing.Dict[str, str]
    parameters: list  # typing.List[Parameter] is not valid in Python

    @property
    def name(self) -> str:
        return self.settings.get("Name", "")

    @property
    def default(self) -> str:
        return self.settings.get("Default", "")

    @classmethod
    def from_dict(cls, in_dict):
        """Initialize from a raw JSON dict"""
        return cls(
            in_dict["Type"],
            in_dict["Flags"],
            in_dict["Settings"],
            in_dict.get("Choices", {}),
            [Parameter.from_dict(x) for x in in_dict.get("Parameters", [])],
        )


ParameterList = typing.List[Parameter]


@dataclass(frozen=True)
class Product:
    """A Smedge product definition"""

    type: UUID
    smedge_class: str
    name: str
    queue: int
    aliases: typing.List[str]
    parameters: ParameterList
    auto_detect: typing.Dict[str, typing.Dict[str, str]]

    @classmethod
    def from_dict(cls, in_dict):
        """Initialize from a raw JSON dict"""
        return cls(
            UUID(in_dict["Type"]),
            in_dict["Class"],
            in_dict["Name"],
            in_dict["Queue"],
            in_dict["Aliases"],
            [Parameter.from_dict(x) for x in in_dict["Parameters"]],
            in_dict["AutoDetect"],
        )

    @cached_property
    def required_parameters(self) -> ParameterList:
        """Gets a list of the parameters that have the "Required" flag set"""
        return sorted(
            [x for x in self.parameters if "Required" in x.flags],
            key=lambda x: x.name,
        )


JobHistoryType = typing.List[HistoryElement]
"""Type used for Job History"""


def find_parameter(
    parameters: ParameterList, name: str
) -> typing.Optional[Parameter]:
    """Finds the given parameter by name regardless of case

    Arguments:
        parameters: a list of Parameter objects
        name: the name to search for (case-insensitive)

    Returns:
        A Parameter object if found, or None
    """
    try:
        return next(x for x in parameters if x.name.lower() == name.lower())
    except StopIteration:
        return None


def conform_job(job: JobType) -> JobType:
    """Ensures that a job's values fit with the requested product

    This will get the product info from the Master, then use it to validate
    the input job object's values. Any lists or dicts will be turned into a
    string based on product settings if possible.

    This is the inverse process of pythonize_job

    Arguments:
        job: The input job dict

    Returns:
        A modified job dict ready to submit

    Raises:
        ValueError: - You did not supply a valid job dict
                    - The job has no type
                    - The type given was not found in the system
                    - Not all required parameters were found
        SmedgeComponentError: Something went wrong
    """

    def recursive_check(
        in_dict: dict, param_list: ParameterList, top_level: bool = False
    ):
        def add_quotes(val: str) -> str:
            val = str(val)
            if (" " in val or "$" in val) and val[0] != '"' and val[-1] != '"':
                val = f'"{val}"'
            return val

        def is_default(val: str) -> bool:
            if val is None:
                return True
            default = param_info.default
            if not default:
                return not val
            if param_info.type in ["Int", "Uint"]:
                return int(val) == int(default)
            if param_info.type == "Float":
                return float(val) == float(default)
            if param_info.type == "Bool":
                return bool(val) == bool(default)
            if param_info.type == "Time":
                if default == "0":
                    return not val or int(val) == 0
                return extract_datetime(val) == extract_datetime(default)
            return val == default

        def check_sequence():
            if param_info.type.endswith("List"):
                separator = param_info.settings.get("Separator", ";")
                enquote = "Enquote" in param_info.flags
                joined_list = separator.join(
                    add_quotes(x) if enquote else str(x) for x in value
                )
                if joined_list and joined_list != param_info.default:
                    result[key] = joined_list
            elif param_info.type == "Multi":
                separator = param_info.settings.get("Separator", " ")
                expecting = param_info.settings.get("Names", "")
                expecting = expecting.split(",")
                if len(value) != len(expecting):
                    logger.warning(
                        f"Sequence {key} has {len(value)} items, but"
                        f" expected {len(expecting)}: {expecting}"
                    )
                else:
                    result[key] = separator.join(str(x) for x in value)
            else:
                logger.warning(
                    f"Sequence in parameter {key} cannot be"
                    f" parsed as a {product_info.type} type."
                )

        def check_mapping():
            if param_info.type == "Multi":
                separator = param_info.settings.get("Separator", " ")
                expecting = param_info.settings.get("Names", "").split(",")
                ordered_values = []
                missing_one = False
                all_empty = True
                for expected_name in expecting:
                    item_value = find_value(value, expected_name)
                    if item_value is None:
                        logger.warning(
                            f"Parameter {key} mapped type missing"
                            f" expected name: {expected_name}"
                        )
                        missing_one = True
                        break
                    if item_value:
                        ordered_values.append(str(item_value))
                        all_empty = False
                if not missing_one and not all_empty:
                    result[key] = separator.join(ordered_values)
            elif param_info.type == "Parameters":
                separator = param_info.settings.get("Separator", " ")
                int_sep = param_info.settings.get("InternalSeparator", " ")
                checked_dict = recursive_check(value, param_info.parameters)
                result_list = [
                    f"{subkey}{int_sep}{subval}"
                    for subkey, subval in checked_dict.items()
                ]
                result[key] = "{" + separator.join(result_list) + "}"
            else:
                logger.warning(
                    f"Mapped parameter {key} cannot be parsed as a"
                    f" {product_info.type} type."
                )

        def check_choice():
            if not value:
                return
            for choice in param_info.choices.values():
                if str(value) == choice:
                    result[key] = choice
                    return
            if str(value) not in param_info.choices.keys():
                logger.warning(f"Parameter {key} has invalid choice: {value}")
            else:
                result[key] = str(value)

        def check_bool_override():
            if value is True:
                result[key] = "1"
            elif value is False:
                result[key] = "0"
            elif value:
                result[key] = str(value)

        result = {}
        for key, value in in_dict.items():
            param_info = find_parameter(param_list, key)
            if not param_info:
                logger.warning(
                    f"Parameter {key} not found in product {product_info.name}"
                )
                continue
            if top_level and param_info.name in required_names:
                required_names.remove(param_info.name)
            if is_default(value):
                continue
            if isinstance(value, (list, tuple)):
                check_sequence()
            elif isinstance(value, dict):
                check_mapping()
            elif param_info.type == "Choice":
                check_choice()
            elif param_info.type == "BoolOverride":
                check_bool_override()
            elif isinstance(value, datetime):
                result[key] = datetime_to_rtime(value)
            elif value is None:
                result[key] = ""
            else:
                result[key] = str(value)
        return result

    product_info = _get_product_info(job)
    required_names = {
        x.name for x in product_info.parameters if "Required" in x.flags
    }
    conformed_job = recursive_check(job, product_info.parameters, True)
    if required_names:
        raise ValueError(f"Missing required parameters: {required_names}")
    return conformed_job


def pythonize_job(job: JobType) -> JobType:
    """Expands values to be more Pythonic.

    Uses the product info to turn IDs into UUID objects, expand lists
    and sub-parameters, and adds all default values from the product.

    This is the inverse process from conform_job.

    Arguments:
        job: The input job dict

    Returns:
        the expanded job dict
    """

    def recursive_pythonize(in_dict: dict, param_list: ParameterList):
        def dequote(val: str, start: str = '"', end: str = '"') -> str:
            if len(val) > 1 and val[0] == start and val[-1] == end:
                val = val[1:-1]
            return val

        result = {}
        for param_info in sorted(param_list, key=lambda x: x.name):
            if param_info.type in ["Separator", "Info", "LastError"]:
                continue
            value = find_value(in_dict, param_info.name)
            if value is None and param_info.type != "BoolOverride":
                value = param_info.default
            if param_info.type.endswith("List"):
                separator = param_info.settings.get("Separator", ";")
                if value:
                    values = value.split(separator)
                    if "Enquote" in param_info.flags:
                        values = [dequote(it) for it in values]
                    result[param_info.name] = values
                else:
                    result[param_info.name] = []
            elif param_info.type == "Multi":
                separator = param_info.settings.get("Separator", " ")
                names = param_info.settings.get("Names", "").split(",")
                values = value.split(separator)
                result[param_info.name] = {
                    names[i]: values[i] if i < len(values) else ""
                    for i in range(len(names))
                }
            elif param_info.type == "Parameters":
                separator = param_info.settings.get("Separator", " ")
                int_sep = param_info.settings.get("InternalSeparator", " ")
                values = dequote(value, "{", "}").split(separator)
                try:
                    result[param_info.name] = recursive_pythonize(
                        {
                            k.strip(): v.strip()
                            for k, v in (it.split(int_sep) for it in values)
                        },
                        param_info.parameters,
                    )
                except ValueError:
                    result[param_info.name] = value
            elif param_info.type == "BoolOverride":
                if value == "1":
                    result[param_info.name] = True
                elif value == "0":
                    result[param_info.name] = False
                else:
                    result[param_info.name] = None
            elif param_info.type == "ID":
                if value and value != "00000000-0000-0000-0000-000000000000":
                    result[param_info.name] = UUID(value)
                else:
                    result[param_info.name] = None
            elif param_info.type == "Time":
                result[param_info.name] = extract_datetime(value)
            elif param_info.type in ["Int", "Uint"]:
                result[param_info.name] = int(value) if value else None
            elif param_info.type == "Float":
                result[param_info.name] = float(value) if value else None
            elif param_info.type == "Bool":
                result[param_info.name] = bool(value) if value else None
            elif not value:
                result[param_info.name] = None
            else:
                result[param_info.name] = value
        return result

    product_info = _get_product_info(job)
    return recursive_pythonize(job, product_info.parameters)


def submit_jobs(
    jobs: typing.Union[JobType, typing.List[JobType]],
    paused: bool = False,
) -> typing.List[UUID]:
    """Submit or update one or more jobs to Smedge

    A Job is a dictionary of values for the job settings. The values to supply
    depend on the type of Job being submitted. See the Smedge Administrator
    Manual for all possible values for all default job types along with
    current defaults and required values. If you do not supply "id" as
    one of the Job dict values, the jobs will get new unique IDs.  Most values
    can be left out to use the defaults.

    If you need to submit a bunch of jobs at the same time, it's far more
    efficient to submit them all at once with a list of jobs than one at a
    time.  If you need to create dependencies between jobs, generate an ID
    for the lead job and use that in the dependent job's WaitForJobID value.

    This will generate a temp file with the job data. If anything goes wrong
    during the submit process, this file will be left in the temp folder
    to allow you to see if job data is causing the problem. If all goes well,
    it will be automatically removed.

    Arguments:
        jobs: a dict supplying the job settings, or a list of these dicts
        paused: set to True to override all jobs to be submitted paused

    Returns:
        A list of Job UUIDs in the same order as the input jobs

    Raises:
        ValueError: - You did not supply a valid job dict or list of dicts
                    - The job's supplied ID is not a valid UUID
                    - More than one job has the same ID
        SmedgeComponentError: Something went wrong
    """
    if not isinstance(jobs, list):
        jobs = [jobs]
    if not jobs:
        raise ValueError("At least one Job must be supplied")
    logger.info(f"Sending {len(jobs)} jobs to Smedge...")
    job_file = NamedTemporaryFile(mode="w", delete=False, suffix=".sj")
    ordered_ids = []
    with job_file as open_file:
        logger.debug(f"Preparing temporary file: {job_file.name}")
        for job in jobs:
            if not isinstance(job, dict):
                raise ValueError("A Job must be a dict")
            if "id" in job:
                job_id = UUID(str(job["id"]))
            elif "ID" in job:
                job_id = UUID(str(job["ID"]))
            elif "Id" in job:
                job_id = UUID(str(job["Id"]))
            elif "iD" in job:
                job_id = UUID(str(job["iD"]))
            else:
                job_id = uuid.uuid4()
            if job_id in ordered_ids:
                raise ValueError(f"Multiple jobs have ID: {job_id}")
            ordered_ids.append(job_id)
            open_file.write(f"[{job_id}]\n")
            for key, value in job.items():
                if key.lower() != "id":
                    if isinstance(value, datetime):
                        value = datetime_to_rtime(value)
                    open_file.write(f"{key} = {value}\n")
    cmd = [
        SUBMIT,
        "FromFile",
        job_file.name,
        "-Update",
    ]
    if paused:
        cmd.append("-Paused")

    result = run_component(cmd)

    expected_ids = ordered_ids.copy()
    for line in result.stdout.split("\n"):
        if line.startswith("Successfully sent new Job ID:"):
            job_id = uuid.UUID(line[-36:])
            if job_id not in expected_ids:
                abort(result, f"Received unknown Job ID: {job_id}")
            expected_ids.remove(job_id)
            logger.info(f"Submitted Job: {job_id}")
    if expected_ids:
        abort(result, f"Master did not acknowledge jobs: {expected_ids}")
    logger.debug(f"Removing temporary file: {job_file.name}")
    os.unlink(job_file.name)
    return ordered_ids


def update_jobs(jobs: NameOrIDList, parameter: str, value: any):
    """Updates a parameter on one or more jobs.

    Arguments:
        jobs: A job ID (string or UUID) or a list of job IDs.
        parameter: The name of the parameter to update.
        value: The new value to set.

    Raises:
        SmedgeComponentError: Something went wrong
    """
    cmd = [SUBMIT, "Update"] + ensure_list(jobs) + [parameter, str(value)]
    run_component(cmd)


def get_product_info(
    product: typing.Optional[NameOrID] = None, use_cache=True
) -> typing.Optional[typing.Union[Product, typing.List[Product]]]:
    """Get product information

    Arguments:
        product: A product to get info about, or None to get all products.
        use_cache: If True, the values are cached and re-used for performance.
                   If False, the cache is discarded and refreshed.

    Returns:
        If you supply a name or ID, will return a dict with the info, or
        None if no product is found with the given name, ID or shortcut.
        If you do not supply a name, will return a list of dicts.

    Raises:
        SmedgeComponentError: Something went wrong
    """
    # Ensure we have a UUID or a lowercase string
    if isinstance(product, str):
        try:
            product = UUID(product)
        except ValueError:
            product = product.lower()
    # Fill the cache if needed
    if not use_cache or not _product_cache:
        cmd = [JOB, "GetProductInfo", "-JSON"]
        result = run_component(cmd)
        if result.stderr:
            abort(result, result.stderr.split("\n")[0])
        else:
            product_list = []
            try:
                product_list = json.loads(result.stdout)
            except json.decoder.JSONDecodeError as e:
                abort(result, str(e))
            if not isinstance(product_list, list):
                abort(
                    result,
                    f"Did not get a list of products: {type(product_list)}",
                )
            _cache_products(product_list)
    # Return requested results if possible
    if product and product in _product_cache:
        return _product_cache[product]
    elif product is None:
        return _product_cache[None]
    else:
        logger.warning(f"No product found to match: {product}")


def ids_to_query(job_ids: NameOrIDList) -> list:
    """Convert an ID or list of IDs (strings or UUIDs) to a Smedge query

    Arguments:
        job_ids: an ID string or UUID object, or a list of these

    Returns:
        a query list suitable for use with get_jobs()

    Raises:
        ValueError: the input list is empty
        TypeError: a supplied value is not a valid job ID
    """
    if not job_ids:
        raise ValueError("No job IDs given")
    if not isinstance(job_ids, list):
        job_ids = [job_ids]
    query = ["any"]
    for job_id in job_ids:
        try:
            if not isinstance(job_id, UUID):
                job_id = UUID(job_id)
            query.append(["ID", "=", job_id])
        except ValueError:
            raise TypeError(f"{job_id} is not a valid job UUID")
    return query


def get_jobs(
    query: typing.Optional[typing.Union[str, UUID, list]] = None,
    get_all_parameters: bool = True,
) -> typing.List[JobType]:
    """Get jobs

    Arguments:
        query: a single job ID as a string or UUID, or a
            JSON string conforming to the Smedge query syntax.
            Gets all jobs with an empty query.
        get_all_parameters: If True, all product parameters are included.
            False returns only non-default values as raw strings.

    Returns:
        A list of dicts for each job that was found. If a Job is not found
        it will not throw an Exception, the job will simply be missing from
        the result.

    Raises:
        SmedgeComponentError: Something went wrong
    """

    class Encoder(json.JSONEncoder):
        """Custom JSON encoder to handle UUIDs"""

        def default(self, obj):
            if isinstance(obj, UUID):
                return str(obj)
            return super().default(obj)

    cmd = [JOB, "Query"]
    if query:
        if not isinstance(query, list):
            query = ids_to_query(query)
        cmd.append(json.dumps(query, cls=Encoder))
    result = run_component(cmd)
    try:
        result_obj = json.loads(result.stdout)
        _cache_products(result_obj["Products"])
        if get_all_parameters:
            return [pythonize_job(job) for job in result_obj["Jobs"]]
        else:
            return result_obj["Jobs"]
    except json.decoder.JSONDecodeError as e:
        abort(result, str(e))


def get_job_history(jobs: NameOrIDList) -> typing.Dict[UUID, JobHistoryType]:
    """Gets the raw job history

    Arguments:
        jobs: one or more job IDs to get the history for

    Returns:
        A list of HistoryElements for every event from the job's history

    Raises:
        SmedgeComponentError: Something went wrong
    """
    history = {}
    if not jobs:
        return history
    cmd = [JOB, "History"]
    cmd.extend(ensure_list(jobs) + ["-Dump"])
    result = run_component(cmd)
    last_id = None
    for line in result.stdout.split("\n"):
        if not line or line == "Name\tWork ID\tTime\tStatus\tEngine ID\tNote":
            continue
        if line.startswith("Job ID: "):
            last_id = UUID(line[8:])
            history[last_id] = []
            continue
        tokens = line.split("\t")
        if len(tokens) < 6:
            abort(result, f"Line does not have all 6 tokens: {line}")
        name = tokens[0]
        work = UUID(tokens[1])
        time = extract_datetime(tokens[2])
        status = int(tokens[3])
        engine = UUID(tokens[4])
        note = tokens[5]
        history[last_id].append(
            HistoryElement(name, work, time, status, engine, note)
        )
    return history


def pause_jobs(jobs: NameOrIDList, stop_work: bool = False):
    """Pause jobs

    Arguments:
        jobs: One or more Job IDs
        stop_work: If True, work from the jobs will also be stopped.
            Otherwise, any running work will be allowed to finish.

    Raises:
        SmedgeComponentError: Something went wrong
    """
    run_simple_command(JOB, "Pause", jobs, "jobs")
    if stop_work:
        run_simple_command(JOB, "Preempt", jobs, "jobs")


def resume_jobs(jobs: NameOrIDList):
    """Resume paused jobs

    Arguments:
        jobs: one or more job IDs

    Raises:
        SmedgeComponentError: Something went wrong
    """
    run_simple_command(JOB, "Resume", jobs, "jobs")


def preempt(jobs: NameOrIDList):
    """Preempt jobs on the farm

    Stops all work from the given jobs without pausing them.
    See also usurp().

    Arguments:
        jobs: one or more job IDs

    Raises:
        SmedgeComponentError: Something went wrong
    """
    run_simple_command(JOB, "Preempt", jobs, "jobs")


def usurp(jobs: NameOrIDList, minimum_priority: typing.Optional[int] = None):
    """Usurp the farm for jobs

    Stops running jobs with lower priority than the given jobs (or supplied
    value). See also preepmpt().

    Arguments:
        jobs: one or more job IDs
        minimum_priority: cut-off priority for usurping the farm (1-100)
            default to the lowest priority job supplied.

    Raises:
        SmedgeComponentError: Something went wrong
    """
    if minimum_priority:
        minimum_priority = str(minimum_priority)
    run_simple_command(JOB, "Usurp", jobs, "jobs", minimum_priority)


def get_job_dispatch_log(jobs: NameOrIDList) -> typing.Dict[UUID, str]:
    """Get dispatch log reports for jobs

    Arguments:
        jobs: one or more job IDs

    Returns:
        dict of the dispatch report for each job mapped to the job ID

    Raises:
        SmedgeComponentError: Something went wrong
    """
    logs = {}
    result = run_simple_command(JOB, "DispatchLog", jobs, "jobs")
    if result:
        last_job = None
        for line in result.stdout.split("\n"):
            if line.startswith("Report for "):
                last_job = uuid.UUID(line[-36:])
                logs[last_job] = []
            elif last_job is None:
                abort(result, f"Expected a Job ID but got: {line}")
            else:
                logs[last_job].append(line)
    return {job: "\n".join(report) for job, report in logs.items()}


def reset_job_failures(jobs: NameOrIDList):
    """Reset the failure counts for the given jobs.

    Arguments:
        jobs: one or more job IDs

    Raises:
        SmedgeComponentError: Something went wrong
    """
    run_simple_command(JOB, "ResetFailures", jobs, "jobs")


def delete_jobs(jobs: NameOrIDList, stop_work: bool = False):
    """Delete jobs

    This does not wait for the jobs to be deleted, and will return
    immediately after successfully sending the request to the Master.

    Arguments:
        jobs: One or more Job IDs
        stop_work: If True, work from the deleted jobs will be terminated.
            Otherwise, work will be allowed to finish, and the job will be
            deleted after all work finishes (no new work will be sent for it).

    Raises:
        SmedgeComponentError: Something went wrong
    """
    run_simple_command(JOB, "Delete", jobs, "jobs", str(stop_work))


def _cache_products(product_list: list):
    _product_cache.clear()
    _product_cache[None] = []
    for product_data in product_list:
        product_obj = Product.from_dict(product_data)
        _product_cache[None].append(product_obj)
        _product_cache[UUID(product_data["Type"])] = product_obj
        _product_cache[product_obj.name.lower()] = product_obj
        for alias in product_obj.aliases:
            _product_cache[alias.lower()] = product_obj


def _get_product_info(job: JobType):
    if not isinstance(job, dict):
        raise ValueError("A job must be a dict")
    product = find_value(job, "Type")
    if product is None:
        raise ValueError("A job must have a 'Type'")
    product_info = get_product_info(product)
    if product_info is None:
        raise ValueError(f"Type was not found: {product}")
    return product_info
