# *************************************************************************
#
#  Copyright (c) 2026 - Datatailr Inc.
#  All Rights Reserved.
#
#  This file is part of Datatailr and subject to the terms and conditions
#  defined in 'LICENSE.txt'. Unauthorized copying and/or distribution
#  of this file, in parts or full, via any medium is strictly prohibited.
# *************************************************************************

from __future__ import annotations

import contextvars
import json
import os
from functools import reduce
from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union
import uuid
import datetime
import re

from datatailr import Image
from datatailr.errors import TaskError
from datatailr.logging import DatatailrLogger
from datatailr.scheduler.base import (
    ACL,
    EntryPoint,
    Environment,
    Job,
    JobType,
    Resources,
    User,
)
from datatailr.scheduler.constants import DEFAULT_TASK_CPU, DEFAULT_TASK_MEMORY
from datatailr.scheduler.arguments_cache import ArgumentsCache
from datatailr.scheduler.schedule import Schedule
from datatailr.utils import (
    is_dt_installed,
    is_remote_cli_dt_installed,
    dict_to_env_vars,
    BiDict,
)

__DAG_CONTEXT__: contextvars.ContextVar = contextvars.ContextVar("dag_context")
logger = DatatailrLogger(os.path.abspath(__file__)).get_logger()

__ARGUMENTS_CACHE__ = None


def get_arguments_cache() -> ArgumentsCache:
    global __ARGUMENTS_CACHE__
    if __ARGUMENTS_CACHE__ is None:
        __ARGUMENTS_CACHE__ = ArgumentsCache()
    return __ARGUMENTS_CACHE__


def get_current_manager():
    return __DAG_CONTEXT__.get(None)


def parse_duration(duration_str: str) -> datetime.timedelta:
    """
    Parses a duration string into a timedelta object.
    Example formats: '1h 30m', '45m', '2h 15s', '10s'
    """
    pattern = r"(?:(\d+)\s*h)?\s*(?:(\d+)\s*m)?\s*(?:(\d+)\s*s)?"
    match = re.match(pattern, duration_str)
    if not match:
        raise ValueError(
            f"Invalid duration string: {duration_str}. Expected format: '<X>h <Y>m <Z>s'"
        )

    hours = int(match.group(1)) if match.group(1) else 0
    minutes = int(match.group(2)) if match.group(2) else 0
    seconds = int(match.group(3)) if match.group(3) else 0

    return datetime.timedelta(hours=hours, minutes=minutes, seconds=seconds)


class CyclicDependencyError(TaskError):
    """
    Exception raised when a cyclic dependency is detected in the batch job dependencies.
    """

    def __init__(self, message: str):
        super().__init__(message)
        self.message = message


class DuplicateJobNameError(TaskError):
    """
    Exception raised when a job with a duplicate name is added to the batch.
    """

    def __init__(self, job_name: str):
        super().__init__(
            f"A job with the name '{job_name}' already exists in the batch."
        )
        self.job_name = job_name


class MissingDagError(TaskError):
    """
    Exception raised when a Task is created outside the context of a Workflow.
    """

    def __init__(self):
        super().__init__(
            "A Task must be either created within the context of a Workflow or a Workflow object has to be provided as the dag argument."
        )


class CodePackageMismatchError(TaskError):
    def __init__(self, message: str):
        super().__init__(message)
        self.message = message


class Task:
    """
    Represents a job within a batch job.

    This class can be extended to define specific configurations for each job in the batch.
    """

    def __init__(
        self,
        name: str,
        entrypoint: EntryPoint,
        resources: Optional[Resources] = None,
        dependencies: Sequence[Union[str, Task]] | None = None,
        dag: Optional[Workflow] = None,
        argument_mapping: Dict[str, str] | None = None,
        env_vars: Dict[str, str | int | float | bool] | None = None,
        id: Optional[int] = None,
    ):
        self.name = name
        self.entrypoint = entrypoint
        self.resources = resources
        self.dependencies: set = set(dependencies or [])
        if dag is None:
            dag = get_current_manager()
        if dag is None:
            raise MissingDagError()
        self.__id = id if id is not None else dag.next_job_id
        self.dag = dag
        self.__args: Dict[str, Any] = {}
        self.dag.__TASK_NAMES__[self.name] = self.__id
        self.dependencies = self.translate_dependencies()
        assert all(
            isinstance(dep, int) for dep in self.dependencies
        ), "All dependencies must be integers representing job IDs."
        inherited_env_vars = self.dag.get_env_vars_copy()
        inherited_env_vars.update(env_vars or {})
        self.__env_vars = inherited_env_vars
        self.dag.add_job(self)
        self.argument_mapping = argument_mapping or {}

    def __call__(self, *args, **kwds) -> Task:
        """
        Allows the Task instance to be called like a function, returning itself.
        This is useful for chaining or functional-style programming.
        """
        return self

    @property
    def args(self) -> Dict[str, Any]:
        """
        Returns the arguments for the Task instance.
        """
        return self.__args or {}

    @args.setter
    def args(self, args: Dict[str, Any]):
        """
        Sets the arguments for the Task instance.
        """
        if not isinstance(args, dict):
            raise TypeError(f"Expected a dictionary for args, got {type(args)}")
        self.__args = args

    @property
    def id(self) -> int:
        """
        Returns the unique identifier of the Task instance.
        """
        return self.__id

    def alias(self, name: str) -> Task:
        """
        Set an alias for the Task instance.

        :param name: The alias name to set.
        """
        if name in self.dag.__TASK_NAMES__:
            raise DuplicateJobNameError(name)
        if self.name in self.dag.__FUNC_CALL_COUNTS__:
            self.dag.__FUNC_CALL_COUNTS__[self.name] -= 1
            if self.dag.__FUNC_CALL_COUNTS__[self.name] == 0:
                self.dag.__FUNC_CALL_COUNTS__.pop(self.name)
            elif self.dag.__FUNC_CALL_COUNTS__[self.name] < 0:
                raise ValueError(f"Invalid function call count for job {self.name}")
        assert self.dag.__TASK_NAMES__.pop(self.name) == self.__id
        self.dag.__TASK_NAMES__[name] = self.__id
        self.name = name
        return self

    def set_resources(
        self,
        resources: Optional[Resources] = None,
        memory: Optional[str] = None,
        cpu: Optional[float] = None,
    ) -> Task:
        """
        Set the resources for the Task instance.

        :param resources: The Resources instance to set.
        """
        if resources is not None:
            if not isinstance(resources, Resources):
                raise TypeError(f"Expected Resources instance, got {type(resources)}")
        else:
            resources = Resources(
                memory=memory or DEFAULT_TASK_MEMORY, cpu=cpu or DEFAULT_TASK_CPU
            )
        self.resources = resources
        return self

    def __repr__(self):
        return (
            f"Task(name={self.name}, entrypoint={self.entrypoint}, "
            f"resources={self.resources}) (id={self.__id})"
        )

    def __getstate__(self) -> object:
        state = self.__dict__.copy()
        state.pop("dag", None)
        return state

    def __setstate__(self, state: dict):
        self.__dict__.update(state)

    def update_env_vars(self, env_vars: Dict[str, str | int | float | bool]) -> Task:
        """
        Update the environment variables for the Task instance.

        :param env_vars: A dictionary of environment variables to update.
        """
        self.__env_vars.update(env_vars)
        return self

    def to_dict(self):
        """
        Convert the Task instance to a dictionary representation.
        """
        return {
            "display_name": self.name,
            "name": self.name,
            "child_number": self.__id,
            "entrypoint": str(self.entrypoint),
            "memory": self.resources.memory if self.resources else DEFAULT_TASK_MEMORY,
            "num_cpus": self.resources.cpu if self.resources else DEFAULT_TASK_CPU,
            "depends_on": list(self.dependencies),
            "env": dict_to_env_vars(self.__env_vars),
        }

    @classmethod
    def from_dict(cls, job_dict: dict, batch: Workflow) -> Task:
        job = cls(
            name=job_dict["display_name"],
            entrypoint=EntryPoint.from_dict(job_dict),
            resources=Resources.from_job_dict(job_dict),
            dependencies=job_dict.get("depends_on", []),
            dag=batch,
            id=job_dict.get("child_number", None),
        )
        return job

    def to_json(self):
        """
        Convert the Task instance to a JSON string representation.
        """
        return json.dumps(self.to_dict())

    def translate_dependencies(self) -> Set[int]:
        """
        Translate the dependencies of the Task instance into a format suitable for the batch job.
        """
        if isinstance(self.dependencies, set) and all(
            isinstance(dep, int) for dep in self.dependencies
        ):
            return self.dependencies

        def get_dependency_name(dep):
            if isinstance(dep, str):
                return dep
            elif isinstance(dep, Task):
                return dep.name
            else:
                raise TypeError(f"Unsupported dependency type: {type(dep)}")

        return set(
            [
                self.dag.__TASK_NAMES__[get_dependency_name(dep)]
                for dep in self.dependencies
            ]
        )

    def __add_dependency__(self, other):
        self.dependencies.add(other.__id)
        arg_name = self.argument_mapping.get(other.name, other.name)
        if arg_name is not None and arg_name in self.__args:
            self.__args[arg_name] = other

    def __lshift__(self, other: Sequence[Task] | Task) -> Sequence[Task] | Task:
        if isinstance(other, list):
            for task in other:
                self.__add_dependency__(task)
        else:
            self.__add_dependency__(other)
        return other

    def __rshift__(self, other: Sequence[Task] | Task) -> Sequence[Task] | Task:
        if isinstance(other, Sequence):
            for task in other:
                if isinstance(task, Task):
                    task.__add_dependency__(self)
        else:
            if isinstance(other, Task):
                other.__add_dependency__(self)
        return other

    def __rrshift__(self, other: Sequence[Task] | Task) -> Task:
        self.__lshift__(other)
        return self

    def __rlshift__(self, other: Sequence[Task] | Task) -> Task:
        self.__rshift__(other)
        return self

    def __hash__(self):
        return self.__id

    def __lt__(self, other: Task) -> bool:
        return self.__id < other.__id

    def run(self):
        """
        Execute the job's entrypoint.
        """
        if isinstance(self.entrypoint, EntryPoint):
            env = {
                "DATATAILR_BATCH_ID": str(self.dag.id),
                "DATATAILR_JOB_ID": str(self.__id),
                "DATATAILR_JOB_NAME": f"{self.dag._internal_name}-{self.__id}",
            }
            self.entrypoint(env=env)
        else:
            raise TypeError(f"Invalid entrypoint type: {type(self.entrypoint)}")


class Workflow(Job):
    """
    Represents a workflow in the scheduler.

    Inherits from Job and is used to define workflows with specific configurations.
    """

    def __init__(
        self,
        name: str,
        environment: Optional[Environment] = Environment.DEV,
        schedule: Optional[Schedule] = None,
        image: Optional[Image] = None,
        run_as: Optional[Union[str, User]] = None,
        resources: Resources = Resources(memory="100m", cpu=1),
        acl: Optional[ACL] = None,
        python_version: str = "auto",
        local_run: bool = False,
        python_requirements: str | List[str] = "",
        build_script_pre: str = "",
        build_script_post: str = "",
        env_vars: Dict[str, str | int | float | bool] | None = None,
        fail_after: datetime.timedelta | str | None = None,
        expire_after: datetime.timedelta | str | None = None,
        get_existing: bool = False,
    ):
        self.__TASK_NAMES__: BiDict = BiDict()
        self.__FUNC_CALL_COUNTS__: Dict[str, int] = {}
        self.__tasks: List[Task] = []
        self.__env_vars = env_vars or {}
        super().__init__(
            environment=environment,
            name=name,
            image=image,
            run_as=run_as,
            resources=resources,
            acl=acl,
            python_version=python_version,
            python_requirements=python_requirements,
            build_script_pre=build_script_pre,
            build_script_post=build_script_post,
            type=JobType.WORKFLOW,
            get_existing=get_existing,
        )
        self._auto_run = False
        self.__next_task_id = -1
        self.__local_run = local_run
        self.__schedule = schedule

        if isinstance(fail_after, str):
            fail_after = parse_duration(fail_after)
        self.__fail_after = fail_after
        if isinstance(expire_after, str):
            expire_after = parse_duration(expire_after)
        self.__expire_after = expire_after

    def reset(self) -> None:
        """
        Reset the Workflow instance to its initial state.
        """
        self.__tasks = []
        self.__next_task_id = -1
        self.__TASK_NAMES__ = BiDict()
        self.__FUNC_CALL_COUNTS__ = {}
        self.__env_vars = {}
        self._auto_run = False

    def set_local_run(self, local_run: bool) -> None:
        """
        Set the local run flag for the Workflow instance.

        :param local_run: A boolean indicating whether to run locally.
        """
        self.__local_run = local_run

    @property
    def next_job_id(self) -> int:
        """
        Returns a generator for the next task ID in the workflow.
        """
        self.__next_task_id += 1
        return self.__next_task_id

    def add_job(self, job: Task):
        """
        Adds a task to the workflow.

        :param job: The Task instance to add.
        """
        if job.name in self.__FUNC_CALL_COUNTS__:
            self.__FUNC_CALL_COUNTS__[job.name] += 1
            job.name = f"{job.name}-{self.__FUNC_CALL_COUNTS__[job.name] - 1}"
        else:
            self.__FUNC_CALL_COUNTS__[job.name] = 1
        # Use the workflow level resource values as defaults for tasks
        job.resources = job.resources or self.resources
        image_path_to_repo = self.image.path_to_repo
        image_path_to_module = self.image.path_to_module
        package_path_to_repo = job.entrypoint.path_to_repo
        package_path_to_module = job.entrypoint.path_to_module

        if image_path_to_repo is None:
            self.image.path_to_repo = package_path_to_repo
        elif package_path_to_repo != image_path_to_repo:
            raise CodePackageMismatchError(
                f"Function {job.entrypoint.function_name} is defined in a different package root: "
                f"{package_path_to_repo} != {image_path_to_repo}"
            )
        if image_path_to_module is None:
            self.image.path_to_module = package_path_to_module
        elif package_path_to_module != image_path_to_module:
            raise CodePackageMismatchError(
                f"Function {job.entrypoint.function_name} is defined in a different module: "
                f"{package_path_to_module} != {image_path_to_module}"
            )
        self.__tasks.append(job)

    def is_task_in(self, task: Task) -> bool:
        return task in self.__tasks

    def get_task_by_name(self, task_name: str) -> Optional[Task]:
        return next((task for task in self.__tasks if task.name == task_name), None)

    def get_task_name_from_id(self, task_id: int) -> Optional[str]:
        return self.__TASK_NAMES__.inverse.get(task_id, None)

    def to_dict(self):
        """
        Convert the Workflow instance to a dictionary representation.
        """
        batch_dict = super().to_dict()
        batch_dict["jobs"] = [job.to_dict() for job in self.__tasks]
        batch_dict["schedule"] = str(self.__schedule) if self.__schedule else None
        batch_dict["fail_after"] = (
            int(
                self.__fail_after.total_seconds() * 1e6 + self.__fail_after.microseconds
            )
            if self.__fail_after
            else None
        )
        batch_dict["expire_after"] = (
            int(
                self.__expire_after.total_seconds() * 1e6
                + self.__expire_after.microseconds
            )
            if self.__expire_after
            else None
        )
        batch_dict.pop("resources", None)
        return batch_dict

    def from_dict(self, job_dict: dict):
        super().from_dict(job_dict)
        self.reset()
        image = self.image
        self.__tasks = [
            Task.from_dict(job_dict | {"image": image}, batch=self)
            for job_dict in job_dict.get("children_full", [])
        ]
        return self

    def to_json(self):
        """
        Convert the Workflow instance to a JSON string representation.
        """
        return json.dumps(self.to_dict(), indent=4)

    def __repr__(self):
        return (
            f"Workflow(name={self.name}, environment={self.environment}, "
            f"run_as={self.run_as}, {len(self.__tasks)} tasks)"
        )

    def set_autorun(self, auto_run):
        self._auto_run = auto_run

    def __enter__(self):
        self._token = __DAG_CONTEXT__.set(self)
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        __DAG_CONTEXT__.reset(self._token)
        self._Job__find_missing_python_packages()
        if self._auto_run:
            self.run()

    def __topological_sort__(self):
        jobs = {
            hash(job): set([hash(dep) for dep in job.dependencies])
            for job in self.__tasks
        }

        for k, v in jobs.items():
            v.discard(k)  # ignore self dependencies
        extra_items_in_deps = reduce(set.union, jobs.values()) - set(jobs.keys())
        jobs.update({item: set() for item in extra_items_in_deps})
        while True:
            ordered = set(item for item, dep in jobs.items() if not dep)
            if not ordered:
                break
            yield sorted(ordered)
            jobs = {
                item: (dep - ordered)
                for item, dep in jobs.items()
                if item not in ordered
            }
        if jobs:
            raise CyclicDependencyError(
                "A cyclic dependency exists amongst {}".format(jobs)
            )

    def get_env_vars_copy(self) -> Dict[str, str | int | float | bool]:
        """
        Returns a copy of the environment variables for the Workflow instance.
        """
        return self.__env_vars.copy()

    def get_schedule_args(self) -> Dict[str, Any]:
        if isinstance(self.__schedule, Schedule):
            args = {
                "at_minutes": self.__schedule.at_minutes,
                "every_minute": self.__schedule.every_minute,
                "at_hours": self.__schedule.at_hours,
                "every_hour": self.__schedule.every_hour,
                "weekdays": self.__schedule.weekdays,
                "day_of_month": self.__schedule.day_of_month,
                "in_month": self.__schedule.in_month,
                "every_month": self.__schedule.every_month,
                "timezone": self.__schedule.timezone,
                "run_after_job_uuid": self.__schedule.run_after_job_uuid,
                "run_after_job_name": self.__schedule.run_after_job_name,
                "run_after_job_condition": self.__schedule.run_after_job_condition,
            }
            args = {key: value for key, value in args.items() if value is not None}
            for key, value in args.items():
                if isinstance(value, list):
                    args[key] = ",".join([str(v) for v in value])
            return args
        return {}

    def prepare_args(self) -> None:
        internal_name = self._internal_name

        def arg_name(arg: Union[Task, str]) -> str:
            return f"{internal_name}-{arg.id}" if isinstance(arg, Task) else arg

        def adjust_mapping(mapping: Dict[str, str]) -> Dict[str, str]:
            result = {}
            for k, v in mapping.items():
                if isinstance(v, Task):
                    result[k] = f"{internal_name}-{v.id}"
                elif isinstance(v, str):
                    job = self.get_task_by_name(v)
                    if job is not None:
                        result[k] = f"{internal_name}-{job.id}"
                    else:
                        result[k] = v
                else:
                    raise TypeError(
                        f"Unsupported type in argument mapping: {type(v)} for key {k}"
                    )
            return result

        def merged(dst: dict[str, str], src: dict[str, str]) -> dict[str, str]:
            out = dict(dst)
            seen_vals = list(out.values())
            for k, v in src.items():
                if v not in seen_vals or not (
                    isinstance(v, str) and v.startswith(internal_name)
                ):
                    out[k] = v
                    seen_vals.append(v)
            return out

        args = {
            f"{internal_name}-{j.id}": merged(
                adjust_mapping(j.argument_mapping),
                {j.argument_mapping.get(k, k): arg_name(v) for k, v in j.args.items()},
            )
            for j in self.__tasks
        }
        get_arguments_cache().add_arguments(self.id, args)

    def save(self) -> Tuple[bool, str]:
        status = super().save()
        if status[0]:
            self.prepare_args()
        return status

    def run(self) -> Tuple[bool, str]:
        if not self.__local_run and (is_dt_installed() or is_remote_cli_dt_installed()):
            result = super().run()
            self.prepare_args()
            return result
        else:
            self.prepare_args()
            os.environ["DATATAILR_BATCH_RUN_ID"] = uuid.uuid4().hex[:8]
            for step in self.__topological_sort__():
                for job_id in step:
                    job = self.__tasks[job_id]
                    logger.info(
                        f"Batch {self.name}, running job '{job.name}' in environment '{self.environment}' as '{self.run_as}'"
                    )
                    job.run()
            from datatailr.scheduler.batch_decorator import __FUNCTIONS_CREATED_IN_DAG__

            __FUNCTIONS_CREATED_IN_DAG__.clear()
            return True, ""
