"""
Implements common methods for any workflow manager
"""
import abc
import time
import logging
from uuid import uuid4
from argparse import Namespace
from collections import defaultdict
from typing import Optional, Generator, Protocol, List, Union, Dict, Tuple

from .script import BaseLoopScript, JobKey, JobDict, Outcome, BaseLoopScriptProtocol


logger = logging.getLogger(__name__)


class WorkFlowManagerProtocol(BaseLoopScriptProtocol, Protocol):
    @abc.abstractmethod
    def get_owned_jobs(self, project_id: Optional[int] = None, **kwargs) -> Generator[JobDict, None, None]:
        ...

    @abc.abstractmethod
    def get_finished_owned_jobs(self, project_id: Optional[int] = None, **kwargs) -> Generator[JobDict, None, None]:
        ...

    @property
    @abc.abstractmethod
    def max_running_jobs(self) -> int:
        ...


class CachedFinishedJobsMixin(WorkFlowManagerProtocol):
    def __init__(self):
        super().__init__()
        self.__finished_cache: Dict[JobKey, Outcome] = {}
        self.__update_finished_cache_called: Dict[int, bool] = defaultdict(bool)

    def update_finished_cache(self, project_id: int):
        if not self.__update_finished_cache_called[project_id]:
            logger.info("Initiating finished cache update.")
            for job in self.get_owned_jobs(project_id, state=["finished"], meta=["close_reason"]):
                if job["key"] in self.__finished_cache:
                    break
                self.__finished_cache[job["key"]] = Outcome(job["close_reason"])
            logger.info(f"Finished jobs cache length: {len(self.__finished_cache)}")
        self.__update_finished_cache_called[project_id] = True

    def get_finished_owned_jobs(self, project_id: Optional[int] = None, **kwargs) -> Generator[JobDict, None, None]:
        kwargs.setdefault("meta", []).append("close_reason")
        # only use for updating finished cache if no filter/limit imposed.
        update_finished_cache = not set(kwargs.keys()).intersection(["count", "has_tag", "lacks_tag"])
        finished_cache = []
        for job in super().get_finished_owned_jobs(project_id, **kwargs):  # type: ignore
            if update_finished_cache:
                finished_cache.append((job["key"], Outcome(job["close_reason"])))
            yield job
        logger.info(f"Preread {len(finished_cache)} finished jobs.")
        while finished_cache:
            key, close_reason = finished_cache.pop()
            self.__finished_cache[key] = close_reason
        logger.info(f"Finished jobs cache length: {len(self.__finished_cache)}")

    def is_finished(self, jobkey: JobKey) -> Optional[Outcome]:
        project_id = int(jobkey.split("/", 1)[0])
        self.update_finished_cache(project_id)
        return self.__finished_cache.get(jobkey)

    def base_loop_tasks(self):
        super().base_loop_tasks()
        for project_id in self.__update_finished_cache_called.keys():
            self.__update_finished_cache_called[project_id] = False


class WorkFlowManager(BaseLoopScript, WorkFlowManagerProtocol):

    # --max-running-job command line option overrides it
    default_max_jobs: int = 1000

    flow_id_required = True

    base_failed_outcomes: Tuple[str, ...] = (
        "failed",
        "killed by oom",
        "cancelled",
        "cancel_timeout",
        "memusage_exceeded",
        "diskusage_exceeded",
        "cancelled (stalled)",
    )

    def __init__(self):
        self.failed_outcomes = list(self.base_failed_outcomes)
        self.is_resumed = False
        self.__autogenerated_flow_id = False
        super().__init__()

    def set_flow_id_name(self, args):
        super().set_flow_id_name(args)
        if not self.name:
            self.argparser.error("Manager name not set.")

    def get_owned_jobs(self, project_id: Optional[int] = None, **kwargs) -> Generator[JobDict, None, None]:
        assert self.flow_id, "This job doesn't have a flow id."
        assert self.name, "This job doesn't have a name."
        assert "has_tag" not in kwargs, "Filtering by flow id requires no extra has_tag."
        assert "state" in kwargs, "'state' parameter must be provided."
        kwargs["has_tag"] = [f"FLOW_ID={self.flow_id}"]
        parent_tag = f"PARENT_NAME={self.name}"
        meta = kwargs.get("meta") or []
        if "tags" not in meta:
            meta.append("tags")
        kwargs["meta"] = meta
        for job in self.get_jobs(project_id, **kwargs):
            if parent_tag in job["tags"]:
                yield job

    def generate_flow_id(self) -> str:
        self.__autogenerated_flow_id = True
        return str(uuid4())

    @property
    def max_running_jobs(self) -> int:
        return self.args.max_running_jobs

    def add_argparser_options(self):
        super().add_argparser_options()
        if not self.name:
            self.argparser.add_argument("name", help="Script name.")
        self.argparser.add_argument(
            "--max-running-jobs",
            type=int,
            default=self.default_max_jobs,
            help="If given, don't allow more than the given jobs running at once.\
                                    Default: %(default)s",
        )

    def parse_args(self) -> Namespace:
        args = super().parse_args()
        if not self.name:
            self.name = args.name
        return args

    def wait_for(
        self,
        jobs_keys: Union[JobKey, List[JobKey]],
        interval: int = 60,
        timeout: float = float("inf"),
        heartbeat: Optional[int] = None,
    ):
        """Waits until all given jobs are not running anymore or until the
        timeout is reached, if a heartbeat is given it'll log an entry every
        heartbeat seconds (considering the interval), otherwise it'll log an
        entry every interval seconds.
        """
        if isinstance(jobs_keys, str):
            jobs_keys = [jobs_keys]
        if heartbeat is None or heartbeat < interval:
            heartbeat = interval
        still_running = dict((key, True) for key in jobs_keys)
        time_waited, next_heartbeat = 0, heartbeat
        while any(still_running.values()) and time_waited < timeout:
            time.sleep(interval)
            time_waited += interval
            for key in (k for k, v in still_running.items() if v):
                if self.is_running(key):
                    if time_waited >= next_heartbeat:
                        next_heartbeat += heartbeat
                        logger.info(f"{key} still running")
                    break
                still_running[key] = False

    def _check_resume_workflow(self):
        for job in self.get_jobs(state=["finished"], meta=["tags"], has_tag=[f"NAME={self.name}"]):
            if self.get_keyvalue_job_tag("FLOW_ID", job["tags"]) == self.flow_id:
                inherited_tags = []
                for tag in job["tags"]:
                    if len(tag.split("=")) == 2:
                        inherited_tags.append(tag)
                self.add_job_tags(tags=inherited_tags)
                self.is_resumed = True
                break

    def get_finished_owned_jobs(self, project_id: Optional[int] = None, **kwargs) -> Generator[JobDict, None, None]:
        kwargs.pop("state", None)
        for job in self.get_owned_jobs(project_id, state=["finished"], **kwargs):
            yield job

    def resume_workflow(self):
        rcount = 0
        for job in self.get_owned_jobs(state=["running", "pending"], meta=["spider_args", "job_cmd", "tags", "spider"]):
            self.resume_running_job_hook(job)
            logger.info(f"Found running job {job['key']}")
            rcount += 1
        if rcount > 0:
            logger.info(f"Found a total of {rcount} running children jobs.")

        fcount = 0
        logger.info("Searching finished children jobs...")
        for job in self.get_finished_owned_jobs(meta=["spider_args", "job_cmd", "tags", "spider", "close_reason"]):
            self.resume_finished_job_hook(job)
            fcount += 1
        if fcount > 0:
            logger.info(f"Found a total of {fcount} finished children jobs.")

    def resume_running_job_hook(self, job: JobDict):
        pass

    def resume_finished_job_hook(self, job: JobDict):
        pass

    def _on_start(self):
        if not self.__autogenerated_flow_id:
            self._check_resume_workflow()
        if self.is_resumed:
            self.resume_workflow()
        super()._on_start()
