import json
from dataclasses import asdict, dataclass
from enum import Enum
from typing import Dict, List, Optional, Type, TypedDict, TypeVar, Union

from dbos._serialization import Serializer
from dbos._sys_db import StepInfo, VersionInfo, WorkflowSchedule, WorkflowStatus


class MessageType(str, Enum):
    EXECUTOR_INFO = "executor_info"
    RECOVERY = "recovery"
    CANCEL = "cancel"
    DELETE = "delete"
    LIST_WORKFLOWS = "list_workflows"
    LIST_QUEUED_WORKFLOWS = "list_queued_workflows"
    RESUME = "resume"
    RESTART = "restart"
    GET_WORKFLOW = "get_workflow"
    EXIST_PENDING_WORKFLOWS = "exist_pending_workflows"
    LIST_STEPS = "list_steps"
    FORK_WORKFLOW = "fork_workflow"
    RETENTION = "retention"
    GET_METRICS = "get_metrics"
    ALERT = "alert"
    EXPORT_WORKFLOW = "export_workflow"
    IMPORT_WORKFLOW = "import_workflow"
    LIST_SCHEDULES = "list_schedules"
    GET_SCHEDULE = "get_schedule"
    PAUSE_SCHEDULE = "pause_schedule"
    RESUME_SCHEDULE = "resume_schedule"
    BACKFILL_SCHEDULE = "backfill_schedule"
    TRIGGER_SCHEDULE = "trigger_schedule"
    LIST_APPLICATION_VERSIONS = "list_application_versions"
    SET_LATEST_APPLICATION_VERSION = "set_latest_application_version"


T = TypeVar("T", bound="BaseMessage")


@dataclass
class BaseMessage:
    type: MessageType
    request_id: str

    @classmethod
    def from_json(cls: Type[T], json_str: str) -> T:
        """
        Safely load a JSON into a dataclass, loading only the
        attributes specified in the dataclass.
        """
        data = json.loads(json_str)
        all_annotations = {}
        for base_cls in cls.__mro__:
            if hasattr(base_cls, "__annotations__"):
                all_annotations.update(base_cls.__annotations__)
        kwargs = {k: v for k, v in data.items() if k in all_annotations}
        return cls(**kwargs)

    def to_json(self) -> str:
        dict_data = asdict(self)
        return json.dumps(dict_data)


@dataclass
class BaseResponse(BaseMessage):
    error_message: Optional[str] = None


@dataclass
class ExecutorInfoRequest(BaseMessage):
    pass


@dataclass
class ExecutorInfoResponse(BaseMessage):
    executor_id: str
    application_version: str
    hostname: Optional[str]
    language: Optional[str]
    dbos_version: Optional[str]
    error_message: Optional[str] = None


@dataclass
class RecoveryRequest(BaseMessage):
    executor_ids: List[str]


@dataclass
class RecoveryResponse(BaseMessage):
    success: bool
    error_message: Optional[str] = None


@dataclass
class CancelRequest(BaseMessage):
    workflow_id: str


@dataclass
class CancelResponse(BaseMessage):
    success: bool
    error_message: Optional[str] = None


@dataclass
class DeleteRequest(BaseMessage):
    workflow_id: str
    delete_children: bool


@dataclass
class DeleteResponse(BaseMessage):
    success: bool
    error_message: Optional[str] = None


@dataclass
class ResumeRequest(BaseMessage):
    workflow_id: str


@dataclass
class ResumeResponse(BaseMessage):
    success: bool
    error_message: Optional[str] = None


@dataclass
class RestartRequest(BaseMessage):
    workflow_id: str


@dataclass
class RestartResponse(BaseMessage):
    success: bool
    error_message: Optional[str] = None


class ListWorkflowsBody(TypedDict, total=False):
    workflow_uuids: List[str]
    workflow_name: Optional[Union[str, List[str]]]
    authenticated_user: Optional[Union[str, List[str]]]
    start_time: Optional[str]
    end_time: Optional[str]
    status: Optional[Union[str, List[str]]]
    application_version: Optional[Union[str, List[str]]]
    forked_from: Optional[Union[str, List[str]]]
    parent_workflow_id: Optional[Union[str, List[str]]]
    queue_name: Optional[Union[str, List[str]]]
    limit: Optional[int]
    offset: Optional[int]
    sort_desc: bool
    workflow_id_prefix: Optional[Union[str, List[str]]]
    load_input: bool
    load_output: bool
    executor_id: Optional[Union[str, List[str]]]
    queues_only: bool


@dataclass
class WorkflowsOutput:
    WorkflowUUID: str
    Status: Optional[str]
    WorkflowName: Optional[str]
    WorkflowClassName: Optional[str]
    WorkflowConfigName: Optional[str]
    AuthenticatedUser: Optional[str]
    AssumedRole: Optional[str]
    AuthenticatedRoles: Optional[str]
    Input: Optional[str]
    Output: Optional[str]
    Error: Optional[str]
    CreatedAt: Optional[str]
    UpdatedAt: Optional[str]
    QueueName: Optional[str]
    ApplicationVersion: Optional[str]
    ExecutorID: Optional[str]
    WorkflowTimeoutMS: Optional[str]
    WorkflowDeadlineEpochMS: Optional[str]
    DeduplicationID: Optional[str]
    Priority: Optional[str]
    QueuePartitionKey: Optional[str]
    ForkedFrom: Optional[str]
    ParentWorkflowID: Optional[str]
    DequeuedAt: Optional[str]

    @classmethod
    def from_workflow_information(cls, info: WorkflowStatus) -> "WorkflowsOutput":
        # Convert fields to strings as needed
        created_at_str = str(info.created_at) if info.created_at is not None else None
        updated_at_str = str(info.updated_at) if info.updated_at is not None else None
        inputs_str = str(info.input) if info.input is not None else None
        outputs_str = str(info.output) if info.output is not None else None
        error_str = str(info.error) if info.error is not None else None
        roles_str = (
            str(info.authenticated_roles)
            if info.authenticated_roles is not None
            else None
        )
        workflow_timeout_ms_str = (
            str(info.workflow_timeout_ms)
            if info.workflow_timeout_ms is not None
            else None
        )
        workflow_deadline_epoch_ms_str = (
            str(info.workflow_deadline_epoch_ms)
            if info.workflow_deadline_epoch_ms is not None
            else None
        )
        priority_str = str(info.priority) if info.priority is not None else None
        dequeued_at_str = (
            str(info.dequeued_at) if info.dequeued_at is not None else None
        )

        return cls(
            WorkflowUUID=info.workflow_id,
            Status=info.status,
            WorkflowName=info.name,
            WorkflowClassName=info.class_name,
            WorkflowConfigName=info.config_name,
            AuthenticatedUser=info.authenticated_user,
            AssumedRole=info.assumed_role,
            AuthenticatedRoles=roles_str,
            Input=inputs_str,
            Output=outputs_str,
            Error=error_str,
            CreatedAt=created_at_str,
            UpdatedAt=updated_at_str,
            QueueName=info.queue_name,
            ApplicationVersion=info.app_version,
            ExecutorID=info.executor_id,
            WorkflowTimeoutMS=workflow_timeout_ms_str,
            WorkflowDeadlineEpochMS=workflow_deadline_epoch_ms_str,
            DeduplicationID=info.deduplication_id,
            Priority=priority_str,
            QueuePartitionKey=info.queue_partition_key,
            ForkedFrom=info.forked_from,
            ParentWorkflowID=info.parent_workflow_id,
            DequeuedAt=dequeued_at_str,
        )


@dataclass
class WorkflowSteps:
    function_id: int
    function_name: str
    output: Optional[str]
    error: Optional[str]
    child_workflow_id: Optional[str]
    started_at_epoch_ms: Optional[str]
    completed_at_epoch_ms: Optional[str]

    @classmethod
    def from_step_info(cls, info: StepInfo) -> "WorkflowSteps":
        output_str = str(info["output"]) if info["output"] is not None else None
        error_str = str(info["error"]) if info["error"] is not None else None
        started_at_str = (
            str(info["started_at_epoch_ms"])
            if info["started_at_epoch_ms"] is not None
            else None
        )
        completed_at_str = (
            str(info["completed_at_epoch_ms"])
            if info["completed_at_epoch_ms"] is not None
            else None
        )
        return cls(
            function_id=info["function_id"],
            function_name=info["function_name"],
            started_at_epoch_ms=started_at_str,
            completed_at_epoch_ms=completed_at_str,
            output=output_str,
            error=error_str,
            child_workflow_id=info["child_workflow_id"],
        )


@dataclass
class ListWorkflowsRequest(BaseMessage):
    body: ListWorkflowsBody


@dataclass
class ListWorkflowsResponse(BaseMessage):
    output: List[WorkflowsOutput]
    error_message: Optional[str] = None


class ListQueuedWorkflowsBody(TypedDict, total=False):
    workflow_uuids: List[str]
    workflow_name: Optional[Union[str, List[str]]]
    authenticated_user: Optional[Union[str, List[str]]]
    start_time: Optional[str]
    end_time: Optional[str]
    status: Optional[Union[str, List[str]]]
    application_version: Optional[Union[str, List[str]]]
    forked_from: Optional[Union[str, List[str]]]
    parent_workflow_id: Optional[Union[str, List[str]]]
    queue_name: Optional[Union[str, List[str]]]
    limit: Optional[int]
    offset: Optional[int]
    sort_desc: bool
    workflow_id_prefix: Optional[Union[str, List[str]]]
    load_input: bool
    load_output: bool
    executor_id: Optional[Union[str, List[str]]]


@dataclass
class ListQueuedWorkflowsRequest(BaseMessage):
    body: ListQueuedWorkflowsBody


@dataclass
class ListQueuedWorkflowsResponse(BaseMessage):
    output: List[WorkflowsOutput]
    error_message: Optional[str] = None


@dataclass
class GetWorkflowRequest(BaseMessage):
    workflow_id: str


@dataclass
class GetWorkflowResponse(BaseMessage):
    output: Optional[WorkflowsOutput]
    error_message: Optional[str] = None


@dataclass
class ExistPendingWorkflowsRequest(BaseMessage):
    executor_id: str
    application_version: str


@dataclass
class ExistPendingWorkflowsResponse(BaseMessage):
    exist: bool
    error_message: Optional[str] = None


@dataclass
class ListStepsRequest(BaseMessage):
    workflow_id: str


@dataclass
class ListStepsResponse(BaseMessage):
    output: Optional[List[WorkflowSteps]]
    error_message: Optional[str] = None


class ForkWorkflowBody(TypedDict):
    workflow_id: str
    start_step: int
    application_version: Optional[str]
    new_workflow_id: Optional[str]


@dataclass
class ForkWorkflowRequest(BaseMessage):
    body: ForkWorkflowBody


@dataclass
class ForkWorkflowResponse(BaseMessage):
    new_workflow_id: Optional[str]
    error_message: Optional[str] = None


class RetentionBody(TypedDict):
    gc_cutoff_epoch_ms: Optional[int]
    gc_rows_threshold: Optional[int]
    timeout_cutoff_epoch_ms: Optional[int]


@dataclass
class RetentionRequest(BaseMessage):
    body: RetentionBody


@dataclass
class RetentionResponse(BaseMessage):
    success: bool
    error_message: Optional[str] = None


@dataclass
class GetMetricsRequest(BaseMessage):
    start_time: str  # ISO 8601
    end_time: str  # ISO 8601
    metric_class: str


@dataclass
class MetricData:
    metric_type: str
    metric_name: str
    value: int


@dataclass
class GetMetricsResponse(BaseMessage):
    metrics: List[MetricData]
    error_message: Optional[str] = None


@dataclass
class AlertRequest(BaseMessage):
    name: str
    message: str
    metadata: Dict[str, str]


@dataclass
class AlertResponse(BaseMessage):
    success: bool
    error_message: Optional[str] = None


@dataclass
class ExportWorkflowRequest(BaseMessage):
    workflow_id: str
    export_children: bool


@dataclass
class ExportWorkflowResponse(BaseMessage):
    serialized_workflow: Optional[str]
    error_message: Optional[str] = None


@dataclass
class ImportWorkflowRequest(BaseMessage):
    serialized_workflow: str


@dataclass
class ImportWorkflowResponse(BaseMessage):
    success: bool
    error_message: Optional[str] = None


# --- Schedule messages ---


@dataclass
class ScheduleOutput:
    schedule_id: str
    schedule_name: str
    workflow_name: str
    workflow_class_name: Optional[str]
    schedule: str
    status: str
    context: str

    @classmethod
    def from_schedule(
        cls, s: WorkflowSchedule, serializer: Serializer
    ) -> "ScheduleOutput":
        context_str = str(serializer.deserialize(s["context"]))
        return cls(
            schedule_id=s["schedule_id"],
            schedule_name=s["schedule_name"],
            workflow_name=s["workflow_name"],
            workflow_class_name=s["workflow_class_name"],
            schedule=s["schedule"],
            status=s["status"],
            context=context_str,
        )


class ListSchedulesBody(TypedDict, total=False):
    status: Optional[Union[str, List[str]]]
    workflow_name: Optional[Union[str, List[str]]]
    schedule_name_prefix: Optional[Union[str, List[str]]]


@dataclass
class ListSchedulesRequest(BaseMessage):
    body: ListSchedulesBody


@dataclass
class ListSchedulesResponse(BaseMessage):
    output: List[ScheduleOutput]
    error_message: Optional[str] = None


@dataclass
class GetScheduleRequest(BaseMessage):
    schedule_name: str


@dataclass
class GetScheduleResponse(BaseMessage):
    output: Optional[ScheduleOutput]
    error_message: Optional[str] = None


@dataclass
class PauseScheduleRequest(BaseMessage):
    schedule_name: str


@dataclass
class PauseScheduleResponse(BaseMessage):
    success: bool
    error_message: Optional[str] = None


@dataclass
class ResumeScheduleRequest(BaseMessage):
    schedule_name: str


@dataclass
class ResumeScheduleResponse(BaseMessage):
    success: bool
    error_message: Optional[str] = None


@dataclass
class BackfillScheduleRequest(BaseMessage):
    schedule_name: str
    start: str  # ISO 8601
    end: str  # ISO 8601


@dataclass
class BackfillScheduleResponse(BaseMessage):
    workflow_ids: List[str]
    error_message: Optional[str] = None


@dataclass
class TriggerScheduleRequest(BaseMessage):
    schedule_name: str


@dataclass
class TriggerScheduleResponse(BaseMessage):
    workflow_id: Optional[str]
    error_message: Optional[str] = None


# --- Application version messages ---


@dataclass
class ApplicationVersionOutput:
    version_id: str
    version_name: str
    version_timestamp: int
    created_at: int

    @classmethod
    def from_version_info(cls, v: VersionInfo) -> "ApplicationVersionOutput":
        return cls(
            version_id=v["version_id"],
            version_name=v["version_name"],
            version_timestamp=v["version_timestamp"],
            created_at=v["created_at"],
        )


@dataclass
class ListApplicationVersionsRequest(BaseMessage):
    pass


@dataclass
class ListApplicationVersionsResponse(BaseMessage):
    output: List[ApplicationVersionOutput]
    error_message: Optional[str] = None


@dataclass
class SetLatestApplicationVersionRequest(BaseMessage):
    version_name: str


@dataclass
class SetLatestApplicationVersionResponse(BaseMessage):
    success: bool
    error_message: Optional[str] = None
