# #region JobLifecycleModule [C:5] [TYPE Module] [SEMANTICS task,lifecycle,state,machine,execution]
# @BRIEF Task creation, execution, pause/resume, and completion transitions for plugin-backed
#        jobs — extracted from the monolithic TaskManager.
# @LAYER Core
# @RELATION DEPENDS_ON -> [TaskGraph]
# @RELATION DEPENDS_ON -> [EventBus]
# @RELATION DEPENDS_ON -> [TaskContext]
# @RELATION DEPENDS_ON -> [PluginLoader]
# @RELATION DEPENDS_ON -> [TaskPersistenceService]
# @RATIONALE Extracted from TaskManager to satisfy INV_7. JobLifecycle encodes the task state
#            machine with plugin execution, context injection, backward-compatible plugin
#            dispatch, and input/mapping pause/resume transitions.
# @REJECTED Keeping lifecycle/execution logic inside TaskManager was rejected — it mixes
#            async execution orchestration with passive registry and log buffering concerns,
#            violating single-responsibility and INV_7.
# @PRE Requested plugin ids resolve to executable plugins; task state transitions target known
#      TaskStatus values.
# @POST Every scheduled task transitions through a valid lifecycle path ending in persisted
#       terminal or waiting state.
# @SIDE_EFFECT Schedules async execution, pauses on input/mapping requests, mutates persisted
#              task lifecycle markers.
# @INVARIANT A task cannot be resumed from a waiting state unless a matching future exists or
#            a new wait future is created.

import asyncio
from concurrent.futures import ThreadPoolExecutor
from datetime import UTC, datetime
import inspect
from typing import Any

from ..cot_logger import seed_trace_id
from ..logger import belief_scope, logger
from .context import TaskContext
from .graph import TaskGraph
from .models import Task, TaskStatus


# #region JobLifecycle [C:5] [TYPE Class] [SEMANTICS task,lifecycle,execution,state,machine]
# @BRIEF Encodes task creation, execution, pause/resume, and completion transitions for
#        plugin-backed jobs.
# @RELATION DEPENDS_ON -> [TaskGraph]
# @RELATION DEPENDS_ON -> [TaskContext]
# @RELATION DEPENDS_ON -> [PluginLoader]
# @RELATION DEPENDS_ON -> [TaskPersistenceService]
# @PRE plugin_loader resolves plugin ids; graph and event_bus are initialized.
# @POST Every scheduled task transitions through a valid lifecycle path ending in persisted
#       terminal or waiting state.
# @RATIONALE Extracted as standalone lifecycle state machine to isolate async execution
#            orchestration (plugin dispatch, context injection, pause/resume futures) from
#            passive registry and log buffering concerns.
# @REJECTED Keeping execution logic inside TaskManager was rejected — it mixed async task
#            orchestration with synchronous registry operations and background threading,
#            violating single-responsibility and making the 708-line module unmaintainable.
class JobLifecycle:
    """Task state machine: creation, execution, pause/resume, completion."""

    # #region __init__ [TYPE Function]
    # @BRIEF Initialize with dependencies.
    def __init__(
        self,
        plugin_loader: Any,
        graph: TaskGraph,
        event_bus: Any,
        persistence_service: Any,
        executor: ThreadPoolExecutor,
        loop: asyncio.AbstractEventLoop,
    ):
        self.plugin_loader = plugin_loader
        self.graph = graph
        self.event_bus = event_bus
        self.persistence_service = persistence_service
        self.executor = executor
        self.loop = loop
        self._dataset_subscribers: dict[str, list[asyncio.Queue]] = {}
    # #endregion __init__

    # #region create_task [C:4] [TYPE Function] [SEMANTICS task,create,persist,schedule]
    # @BRIEF Creates and queues a new task for execution.
    # @PRE Plugin with plugin_id exists. Params are valid dict.
    # @POST Task is created, added to registry, and scheduled for execution.
    # @RAISES ValueError if plugin not found or params invalid.
    async def create_task(
        self, plugin_id: str, params: dict[str, Any], user_id: str | None = None,
        add_log_callback=None,
    ) -> Task:
        with belief_scope("JobLifecycle.create_task", f"plugin_id={plugin_id}"):
            if not self.plugin_loader.has_plugin(plugin_id):
                logger.error(f"Plugin with ID '{plugin_id}' not found.")
                raise ValueError(f"Plugin with ID '{plugin_id}' not found.")
            self.plugin_loader.get_plugin(plugin_id)
            if not isinstance(params, dict):
                logger.error("Task parameters must be a dictionary.")
                raise ValueError("Task parameters must be a dictionary.")
            logger.reason("Creating task node and scheduling execution")
            task = Task(plugin_id=plugin_id, params=params, user_id=user_id)
            self.graph.add_task(task)
            self.persistence_service.persist_task(task)
            logger.info(f"Task {task.id} created and scheduled for execution")
            self.loop.create_task(
                self._run_task(task.id, add_log_callback=add_log_callback)
            )
            logger.reflect(
                "Task creation persisted and execution scheduled",
                extra={"task_id": task.id, "plugin_id": plugin_id},
            )
            return task
    # #endregion create_task

    # #region _run_task [C:4] [TYPE Function] [SEMANTICS task,execute,context,dispatch]
    # @BRIEF Internal method to execute a task with TaskContext support.
    # @PRE Task exists in registry.
    # @POST Task is executed, status updated to SUCCESS or FAILED.
    # @SIDE_EFFECT Executes plugin code, updates task status and persistence, emits logs.
    async def _run_task(self, task_id: str, add_log_callback=None):
        seed_trace_id()
        with belief_scope("JobLifecycle._run_task", f"task_id={task_id}"):
            task = self.graph.get_task(task_id)
            if task is None:
                logger.error(f"Task {task_id} not found in registry")
                return
            plugin = self.plugin_loader.get_plugin(task.plugin_id)
            logger.reason(
                "Transitioning task to running state",
                extra={"task_id": task_id, "plugin_id": task.plugin_id},
            )
            logger.info(
                f"Starting execution of task {task_id} for plugin '{plugin.name}'"
            )
            task.status = TaskStatus.RUNNING
            task.started_at = datetime.now(UTC)
            self.persistence_service.persist_task(task)
            if add_log_callback:
                add_log_callback(
                    task_id, "INFO",
                    f"Task started for plugin '{plugin.name}'",
                    source="system",
                )
            try:
                params = {**task.params, "_task_id": task_id}
                sig = inspect.signature(plugin.execute)
                accepts_context = "context" in sig.parameters
                if accepts_context:
                    # Create TaskContext for new-style plugins
                    context = TaskContext(
                        task_id=task_id,
                        add_log_fn=add_log_callback,
                        params=params,
                        default_source="plugin",
                        background_tasks=None,
                    )
                    if asyncio.iscoroutinefunction(plugin.execute):
                        task.result = await plugin.execute(params, context=context)
                    else:
                        task.result = await self.loop.run_in_executor(
                            self.executor,
                            lambda: plugin.execute(params, context=context),
                        )
                else:
                    if asyncio.iscoroutinefunction(plugin.execute):
                        task.result = await plugin.execute(params)
                    else:
                        task.result = await self.loop.run_in_executor(
                            self.executor, plugin.execute, params
                        )
                logger.info(f"Task {task_id} completed successfully")
                task.status = TaskStatus.SUCCESS
                if add_log_callback:
                    add_log_callback(
                        task_id, "INFO",
                        f"Task completed successfully for plugin '{plugin.name}'",
                        source="system",
                    )
            except Exception as e:
                logger.error(f"Task {task_id} failed: {e}")
                task.status = TaskStatus.FAILED
                if add_log_callback:
                    add_log_callback(
                        task_id, "ERROR",
                        f"Task failed: {e}",
                        source="system",
                        metadata={"error_type": type(e).__name__},
                    )
            finally:
                task.finished_at = datetime.now(UTC)
                self.event_bus.flush_task_logs(task_id)
                self.persistence_service.persist_task(task)
                logger.info(
                    f"Task {task_id} execution finished with status: {task.status}"
                )
                logger.reflect(
                    "Task lifecycle reached persisted terminal state",
                    extra={"task_id": task_id, "status": str(task.status)},
                )
                # Broadcast dataset.updated for mapping and documentation tasks (FR-024, R4)
                if task.plugin_id in ("dataset-mapper", "llm_documentation"):
                    dataset_ids = task.params.get("dataset_ids") or [task.params.get("dataset_id")]
                    env_id = task.params.get("env") or task.params.get("environment_id", "")
                    if dataset_ids and env_id:
                        self._broadcast_dataset_updated(env_id, dataset_ids)
    # #endregion _run_task

    # #region resolve_task [C:3] [TYPE Function] [SEMANTICS task,resolve,mapping,resume]
    # @BRIEF Resumes a task that is awaiting mapping.
    # @PRE Task exists and is in AWAITING_MAPPING state.
    # @POST Task status updated to RUNNING, params updated, execution resumed.
    # @RAISES ValueError if task not found or not awaiting mapping.
    async def resolve_task(
        self, task_id: str, resolution_params: dict[str, Any],
    ) -> None:
        with belief_scope("JobLifecycle.resolve_task", f"task_id={task_id}"):
            task = self.graph.get_task(task_id)
            if not task or task.status != TaskStatus.AWAITING_MAPPING:
                raise ValueError("Task is not awaiting mapping.")
            task.params.update(resolution_params)
            task.status = TaskStatus.RUNNING
            self.persistence_service.persist_task(task)
            self.graph.resolve_future(task_id, True)
    # #endregion resolve_task

    # #region wait_for_resolution [C:3] [TYPE Function] [SEMANTICS task,wait,mapping,future]
    # @BRIEF Pauses execution and waits for a resolution signal.
    # @PRE Task exists.
    # @POST Execution pauses until future is set.
    async def wait_for_resolution(self, task_id: str) -> None:
        with belief_scope("JobLifecycle.wait_for_resolution", f"task_id={task_id}"):
            task = self.graph.get_task(task_id)
            if not task:
                return
            task.status = TaskStatus.AWAITING_MAPPING
            self.persistence_service.persist_task(task)
            future = self.loop.create_future()
            self.graph.create_future(task_id, future)
            try:
                await future
            finally:
                self.graph.remove_future(task_id)
    # #endregion wait_for_resolution

    # #region wait_for_input [C:3] [TYPE Function] [SEMANTICS task,wait,input,future]
    # @BRIEF Pauses execution and waits for user input.
    # @PRE Task exists.
    # @POST Execution pauses until future is set via resume_task_with_password.
    async def wait_for_input(self, task_id: str) -> None:
        with belief_scope("JobLifecycle.wait_for_input", f"task_id={task_id}"):
            task = self.graph.get_task(task_id)
            if not task:
                return
            future = self.loop.create_future()
            self.graph.create_future(task_id, future)
            try:
                await future
            finally:
                self.graph.remove_future(task_id)
    # #endregion wait_for_input

    # #region await_input [C:3] [TYPE Function] [SEMANTICS task,input,pause,state]
    # @BRIEF Transition a task to AWAITING_INPUT state with input request.
    # @PRE Task exists and is in RUNNING state.
    # @POST Task status changed to AWAITING_INPUT, input_request set, persisted.
    # @RAISES ValueError if task not found or not RUNNING.
    def await_input(
        self, task_id: str, input_request: dict[str, Any],
        add_log_callback=None,
    ) -> None:
        with belief_scope("JobLifecycle.await_input", f"task_id={task_id}"):
            task = self.graph.get_task(task_id)
            if not task:
                raise ValueError(f"Task {task_id} not found")
            if task.status != TaskStatus.RUNNING:
                raise ValueError(
                    f"Task {task_id} is not RUNNING (current: {task.status})"
                )
            task.status = TaskStatus.AWAITING_INPUT
            task.input_required = True
            task.input_request = input_request
            self.persistence_service.persist_task(task)
            if add_log_callback:
                add_log_callback(
                    task_id, "INFO",
                    "Task paused for user input",
                    metadata={"input_request": input_request},
                )
    # #endregion await_input

    # #region resume_task_with_password [C:3] [TYPE Function] [SEMANTICS task,resume,password,input]
    # @BRIEF Resume a task that is awaiting input with provided passwords.
    # @PRE Task exists and is in AWAITING_INPUT state.
    # @POST Task status changed to RUNNING, passwords injected, task resumed.
    # @RAISES ValueError if task not found, not awaiting input, or passwords invalid.
    def resume_task_with_password(
        self, task_id: str, passwords: dict[str, str],
        add_log_callback=None,
    ) -> None:
        with belief_scope(
            "JobLifecycle.resume_task_with_password", f"task_id={task_id}"
        ):
            task = self.graph.get_task(task_id)
            if not task:
                raise ValueError(f"Task {task_id} not found")
            if task.status != TaskStatus.AWAITING_INPUT:
                raise ValueError(
                    f"Task {task_id} is not AWAITING_INPUT (current: {task.status})"
                )
            if not isinstance(passwords, dict) or not passwords:
                raise ValueError("Passwords must be a non-empty dictionary")
            task.params["passwords"] = passwords
            task.input_required = False
            task.input_request = None
            task.status = TaskStatus.RUNNING
            self.persistence_service.persist_task(task)
            if add_log_callback:
                add_log_callback(
                    task_id, "INFO",
                    "Task resumed with passwords",
                    metadata={"databases": list(passwords.keys())},
                )
            self.graph.resolve_future(task_id, True)
    # #endregion resume_task_with_password

    # #region subscribe_dataset_events [C:2] [TYPE Function] [SEMANTICS dataset,subscribe,events,websocket]
    # @BRIEF Subscribe to dataset.updated events for a specific environment.
    async def subscribe_dataset_events(self, env_id: str) -> asyncio.Queue:
        queue = asyncio.Queue()
        if env_id not in self._dataset_subscribers:
            self._dataset_subscribers[env_id] = []
        self._dataset_subscribers[env_id].append(queue)
        return queue
    # #endregion subscribe_dataset_events

    # #region unsubscribe_dataset_events [C:2] [TYPE Function] [SEMANTICS dataset,unsubscribe,events]
    # @BRIEF Unsubscribe from dataset.updated events for a specific environment.
    def unsubscribe_dataset_events(self, env_id: str, queue: asyncio.Queue):
        if env_id in self._dataset_subscribers:
            if queue in self._dataset_subscribers[env_id]:
                self._dataset_subscribers[env_id].remove(queue)
            if not self._dataset_subscribers[env_id]:
                del self._dataset_subscribers[env_id]
    # #endregion unsubscribe_dataset_events

    # #region _broadcast_dataset_updated [C:3] [TYPE Function] [SEMANTICS dataset,broadcast,event,websocket]
    # @BRIEF Broadcast dataset.updated event to all subscribers for a given environment.
    def _broadcast_dataset_updated(self, env_id: str, dataset_ids: list[int]):
        event = {"type": "dataset.updated", "payload": {"env_id": env_id, "dataset_ids": dataset_ids}}
        if env_id in self._dataset_subscribers:
            for queue in self._dataset_subscribers[env_id]:
                try:
                    self.loop.call_soon_threadsafe(queue.put_nowait, event)
                except Exception:
                    pass
    # #endregion _broadcast_dataset_updated
# #endregion JobLifecycle
# #endregion JobLifecycleModule