from __future__ import annotations

from collections.abc import Callable
from datetime import timedelta
from typing import Any

from temporalio.client import Client, WorkflowExecutionStatus
from temporalio.common import RetryPolicy as TemporalRetryPolicy
from temporalio.common import WorkflowIDReusePolicy
from temporalio.service import RPCError, RPCStatusCode

from terminaluse.lib.core.clients.temporal.types import (
    DuplicateWorkflowPolicy,
    RetryPolicy,
    TaskStatus,
    WorkflowState,
)
from terminaluse.lib.core.clients.temporal.utils import get_temporal_client
from terminaluse.lib.utils.logging import make_logger
from terminaluse.lib.utils.model_utils import BaseModel

logger = make_logger(__name__)

DEFAULT_RETRY_POLICY = RetryPolicy(
    maximum_attempts=1,
    initial_interval=timedelta(seconds=1),
    backoff_coefficient=2.0,
    maximum_interval=timedelta(minutes=10),
)


TEMPORAL_STATUS_TO_UPLOAD_STATUS_AND_REASON = {
    WorkflowExecutionStatus.CANCELED: WorkflowState(
        status=TaskStatus.CANCELED,
        reason="Task canceled by the user.",
        is_terminal=True,
    ),
    WorkflowExecutionStatus.COMPLETED: WorkflowState(
        status=TaskStatus.COMPLETED,
        reason="Task completed successfully.",
        is_terminal=True,
    ),
    WorkflowExecutionStatus.FAILED: WorkflowState(
        status=TaskStatus.FAILED,
        reason="Task encountered terminal failure. Please contact support if retrying does not resolve the issue.",
        is_terminal=True,
    ),
    WorkflowExecutionStatus.RUNNING: WorkflowState(
        status=TaskStatus.RUNNING,
        reason="Task is running.",
        is_terminal=False,
    ),
    WorkflowExecutionStatus.TERMINATED: WorkflowState(
        status=TaskStatus.CANCELED,
        reason="Task canceled by the user.",
        is_terminal=True,
    ),
    WorkflowExecutionStatus.TIMED_OUT: WorkflowState(
        status=TaskStatus.FAILED,
        reason="Task timed out. Please contact support if retrying does not resolve the issue",
        is_terminal=True,
    ),
    WorkflowExecutionStatus.CONTINUED_AS_NEW: WorkflowState(
        status=TaskStatus.RUNNING,
        reason="Task is running.",
        is_terminal=False,
    ),
}

DUPLICATE_POLICY_TO_ID_REUSE_POLICY = {
    DuplicateWorkflowPolicy.ALLOW_DUPLICATE: WorkflowIDReusePolicy.ALLOW_DUPLICATE,
    DuplicateWorkflowPolicy.ALLOW_DUPLICATE_FAILED_ONLY: WorkflowIDReusePolicy.ALLOW_DUPLICATE_FAILED_ONLY,
    DuplicateWorkflowPolicy.REJECT_DUPLICATE: WorkflowIDReusePolicy.REJECT_DUPLICATE,
    DuplicateWorkflowPolicy.TERMINATE_IF_RUNNING: WorkflowIDReusePolicy.TERMINATE_IF_RUNNING,
}


class TemporalClient:
    def __init__(self, temporal_client: Client | None = None, plugins: list[Any] = []):
        self._client: Client | None = temporal_client
        self._plugins = plugins

    @property
    def client(self) -> Client:
        """Get the temporal client, raising an error if not initialized."""
        if self._client is None:
            raise RuntimeError("Temporal client not initialized - ensure temporal_address is properly configured")
        return self._client

    @classmethod
    async def create(cls, temporal_address: str, plugins: list[Any] = []):
        if temporal_address in [
            "false",
            "False",
            "null",
            "None",
            "",
            "undefined",
            False,
            None,
        ]:
            _client = None
        else:
            _client = await get_temporal_client(temporal_address, plugins=plugins)
        return cls(_client, plugins)

    async def setup(self, temporal_address: str):
        self._client = await self._get_temporal_client(temporal_address=temporal_address)

    async def _get_temporal_client(self, temporal_address: str) -> Client | None:
        if temporal_address in [
            "false",
            "False",
            "null",
            "None",
            "",
            "undefined",
            False,
            None,
        ]:
            return None
        else:
            return await get_temporal_client(temporal_address, plugins=self._plugins)

    async def start_workflow(
        self,
        *args: Any,
        duplicate_policy: DuplicateWorkflowPolicy = DuplicateWorkflowPolicy.ALLOW_DUPLICATE,
        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
        task_timeout: timedelta = timedelta(seconds=10),
        execution_timeout: timedelta = timedelta(seconds=86400),
        **kwargs: Any,
    ) -> str:
        temporal_retry_policy = TemporalRetryPolicy(**retry_policy.model_dump(exclude_unset=True))
        workflow_handle = await self.client.start_workflow(
            *args,
            retry_policy=temporal_retry_policy,
            task_timeout=task_timeout,
            execution_timeout=execution_timeout,
            id_reuse_policy=DUPLICATE_POLICY_TO_ID_REUSE_POLICY[duplicate_policy],
            **kwargs,
        )
        return workflow_handle.id

    async def send_signal(
        self,
        workflow_id: str,
        signal: str | Callable[[dict[str, Any] | list[Any] | str | int | float | bool | BaseModel], Any],
        payload: dict[str, Any] | list[Any] | str | int | float | bool | BaseModel,
    ) -> None:
        handle = self.client.get_workflow_handle(workflow_id=workflow_id)
        await handle.signal(signal, payload)  # type: ignore[arg-type]

    async def query_workflow(
        self,
        workflow_id: str,
        query: str | Callable[[dict[str, Any] | list[Any] | str | int | float | bool | BaseModel], Any],
    ) -> Any:
        """
        Submit a query to a workflow by name and return the results.

        Args:
            workflow_id: The ID of the workflow to query
            query: The name of the query or a callable query function

        Returns:
            The result of the query
        """
        handle = self.client.get_workflow_handle(workflow_id=workflow_id)
        return await handle.query(query)

    async def get_workflow_status(self, workflow_id: str) -> WorkflowState:
        try:
            handle = self.client.get_workflow_handle(workflow_id=workflow_id)
            description = await handle.describe()
            status = description.status
            if status is None:
                return WorkflowState(
                    status="UNKNOWN",
                    reason="Workflow status is unknown",
                    is_terminal=False,
                )
            return TEMPORAL_STATUS_TO_UPLOAD_STATUS_AND_REASON[status]
        except RPCError as e:
            if e.status == RPCStatusCode.NOT_FOUND:
                return WorkflowState(
                    status="NOT_FOUND",
                    reason="Workflow not found",
                    is_terminal=True,
                )
            raise

    async def terminate_workflow(self, workflow_id: str) -> None:
        return await self.client.get_workflow_handle(workflow_id).terminate()

    async def cancel_workflow(self, workflow_id: str) -> None:
        return await self.client.get_workflow_handle(workflow_id).cancel()
