import threading
from dataclasses import dataclass
from datetime import datetime
from typing import Literal, TypedDict

from sqlalchemy.orm import Session, sessionmaker

import contexts
from core.workflow.node_execution_process_data import (
    WORKFLOW_TOOL_INVOCATION_ID_KEY,
    WORKFLOW_TOOL_ROOT_APP_ID_KEY,
)
from core.workflow.nodes.human_input.pause_reason import HumanInputRequired
from graphon.enums import WorkflowExecutionStatus
from libs.infinite_scroll_pagination import InfiniteScrollPagination
from machinery.context import RequestContext
from models import (
    WorkflowRun,
    WorkflowRunTriggeredFrom,
)
from repositories.api_workflow_node_execution_repository import DifyAPIWorkflowNodeExecutionRepository
from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository
from services.workflow_node_execution_trace_service import (
    WorkflowNodeExecutionTrace,
    assemble_workflow_node_execution_traces,
)
from services.workflow_pause_service import load_workflow_pause_snapshot


class WorkflowRunListArgs(TypedDict, total=False):
    """Expected shape of the args dict passed to workflow run pagination methods."""

    limit: int
    last_id: str
    status: str


@dataclass(frozen=True, slots=True)
class WorkflowRunPausedNode:
    node_id: str
    node_title: str
    form_id: str
    form_version: Literal["1", "2"]
    form_token: str | None


@dataclass(frozen=True, slots=True)
class WorkflowRunPauseDetails:
    paused_at: datetime | None
    paused_nodes: tuple[WorkflowRunPausedNode, ...]


class WorkflowRunService:
    def __init__(
        self,
        *,
        workflow_runs: DifyAPISQLAlchemyWorkflowRunRepository,
        node_executions: DifyAPIWorkflowNodeExecutionRepository,
        session_factory: sessionmaker[Session],
    ) -> None:
        self._workflow_runs = workflow_runs
        self._node_executions = node_executions
        self._session_factory = session_factory

    def get_paginate_advanced_chat_workflow_runs(
        self,
        context: RequestContext,
        *,
        app_id: str,
        args: WorkflowRunListArgs,
        triggered_from: WorkflowRunTriggeredFrom = WorkflowRunTriggeredFrom.DEBUGGING,
    ) -> InfiniteScrollPagination:
        """
        Get advanced chat app workflow run list

        :param context: admitted Console request context
        :param app_id: app id
        :param args: request args
        :param triggered_from: workflow run triggered from (default: DEBUGGING for preview runs)
        """

        class WorkflowWithMessage:
            message_id: str
            conversation_id: str

            def __init__(self, workflow_run: WorkflowRun):
                self._workflow_run = workflow_run

            def __getattr__(self, item):
                return getattr(self._workflow_run, item)

        pagination = self.get_paginate_workflow_runs(
            context,
            app_id=app_id,
            args=args,
            triggered_from=triggered_from,
        )

        # Batch-load the associated Message for every run in a single query to avoid
        # an N+1 pattern: the deprecated WorkflowRun.message property issues one query
        # per run. The filter matches that property exactly (app_id + workflow_run_id).
        workflow_runs = pagination.data
        run_ids = [workflow_run.id for workflow_run in workflow_runs]
        messages_by_run_id = self._workflow_runs.get_message_refs(
            app_id=app_id,
            workflow_run_ids=run_ids,
        )

        with_message_workflow_runs = []
        for workflow_run in workflow_runs:
            message = messages_by_run_id.get(workflow_run.id)
            with_message_workflow_run = WorkflowWithMessage(workflow_run=workflow_run)
            if message:
                with_message_workflow_run.message_id = message.message_id
                with_message_workflow_run.conversation_id = message.conversation_id

            with_message_workflow_runs.append(with_message_workflow_run)

        pagination.data = with_message_workflow_runs
        return pagination

    def get_paginate_workflow_runs(
        self,
        context: RequestContext,
        *,
        app_id: str,
        args: WorkflowRunListArgs,
        triggered_from: WorkflowRunTriggeredFrom = WorkflowRunTriggeredFrom.DEBUGGING,
    ) -> InfiniteScrollPagination:
        """
        Get workflow run list

        :param context: admitted Console request context
        :param app_id: app id
        :param args: request args
        :param triggered_from: workflow run triggered from (default: DEBUGGING)
        """
        limit = int(args.get("limit", 20))
        last_id = args.get("last_id")
        status = args.get("status")

        return self._workflow_runs.get_paginated_workflow_runs(
            tenant_id=context.active_workspace_id,
            app_id=app_id,
            triggered_from=triggered_from,
            limit=limit,
            last_id=last_id,
            status=status,
        )

    def get_workflow_run(self, context: RequestContext, *, app_id: str, run_id: str) -> WorkflowRun | None:
        """
        Get workflow run detail

        :param context: admitted Console request context
        :param app_id: app id
        :param run_id: workflow run id
        """
        return self._workflow_runs.get_workflow_run_by_id(
            tenant_id=context.active_workspace_id,
            app_id=app_id,
            run_id=run_id,
        )

    def get_workflow_runs_count(
        self,
        context: RequestContext,
        *,
        app_id: str,
        status: str | None = None,
        time_range: str | None = None,
        triggered_from: WorkflowRunTriggeredFrom = WorkflowRunTriggeredFrom.DEBUGGING,
    ) -> dict[str, int]:
        """
        Get workflow runs count statistics

        :param context: admitted Console request context
        :param app_id: app id
        :param status: optional status filter
        :param time_range: optional time range filter (e.g., "7d", "4h", "30m", "30s")
        :param triggered_from: workflow run triggered from (default: DEBUGGING)
        :return: dict with total and status counts
        """
        return self._workflow_runs.get_workflow_runs_count(
            tenant_id=context.active_workspace_id,
            app_id=app_id,
            triggered_from=triggered_from,
            status=status,
            time_range=time_range,
        )

    def get_workflow_run_node_executions(
        self,
        context: RequestContext,
        *,
        app_id: str,
        run_id: str,
    ) -> list[WorkflowNodeExecutionTrace]:
        """
        Get workflow run node execution list

        Returns an empty list when the run is not visible to the admitted workspace.
        Creator accounts and end users are resolved in a short read-only session that
        closes before returning, so the traces carry detached, already-loaded rows.
        """
        workflow_run = self.get_workflow_run(context, app_id=app_id, run_id=run_id)

        contexts.plugin_tool_providers.set({})
        contexts.plugin_tool_providers_lock.set(threading.Lock())

        if not workflow_run:
            return []

        node_executions = self._node_executions.get_executions_by_workflow_run(
            tenant_id=context.active_workspace_id,
            app_id=app_id,
            workflow_run_id=run_id,
        )
        with self._session_factory() as session:
            return assemble_workflow_node_execution_traces(node_executions, self._node_executions, session=session)

    def get_workflow_tool_node_executions(
        self,
        context: RequestContext,
        *,
        app_id: str,
        run_id: str,
        node_execution_id: str,
    ) -> list[WorkflowNodeExecutionTrace] | None:
        """Read nested Tool traces through the admitted root app and workflow run.

        Returns None when the root run is not visible to the admitted workspace.
        Creator lookups use a short read-only session that closes before the internal
        correlation keys are stripped from ``process_data``.
        """
        if self.get_workflow_run(context, app_id=app_id, run_id=run_id) is None:
            return None
        contexts.plugin_tool_providers.set({})
        contexts.plugin_tool_providers_lock.set(threading.Lock())
        executions = self._node_executions.get_workflow_tool_executions(
            tenant_id=context.active_workspace_id,
            workflow_run_id=run_id,
            parent_node_execution_id=node_execution_id,
        )
        with self._session_factory() as session:
            traces = assemble_workflow_node_execution_traces(executions, self._node_executions, session=session)
        for trace in traces:
            if trace.process_data is not None:
                trace.process_data = {
                    key: value
                    for key, value in trace.process_data.items()
                    if key
                    not in (
                        WORKFLOW_TOOL_INVOCATION_ID_KEY,
                        WORKFLOW_TOOL_ROOT_APP_ID_KEY,
                    )
                }
        return traces

    def get_pause_details(
        self,
        context: RequestContext,
        *,
        workflow_run_id: str,
    ) -> WorkflowRunPauseDetails | None:
        with self._session_factory() as session:
            pause_record = self._workflow_runs.get_pause_record(
                session,
                workspace_id=context.active_workspace_id,
                workflow_run_id=workflow_run_id,
            )
        if pause_record is None:
            return None
        if pause_record.status != WorkflowExecutionStatus.PAUSED or pause_record.pause is None:
            return WorkflowRunPauseDetails(paused_at=None, paused_nodes=())

        snapshot = load_workflow_pause_snapshot(pause_record.pause, session_factory=self._session_factory)
        human_input_reasons: list[HumanInputRequired] = []
        for reason in snapshot.reasons:
            if not isinstance(reason, HumanInputRequired):
                raise NotImplementedError(f"Pause details do not support {type(reason).__name__}")
            human_input_reasons.append(reason)

        legacy_form_ids = [reason.form_id for reason in human_input_reasons if reason.form_version == "1"]
        form_tokens: dict[str, str] = {}
        if legacy_form_ids:
            with self._session_factory() as session:
                form_tokens = self._workflow_runs.get_legacy_form_tokens(session, form_ids=legacy_form_ids)
        form_tokens.update(
            {
                form_id: prepared.form_token.get_secret_value()
                for form_id, prepared in snapshot.v2_forms.items()
                if prepared.form_token is not None
            }
        )

        return WorkflowRunPauseDetails(
            paused_at=pause_record.pause.paused_at,
            paused_nodes=tuple(
                WorkflowRunPausedNode(
                    node_id=reason.node_id,
                    node_title=reason.node_title,
                    form_id=reason.form_id,
                    form_version=reason.form_version,
                    form_token=form_tokens.get(reason.form_id),
                )
                for reason in human_input_reasons
            ),
        )