from __future__ import annotations

import asyncio
from datetime import UTC, datetime
from typing import Any, AsyncIterator

from airflow.sdk import BaseOperator
from airflow.triggers.base import BaseTrigger, TriggerEvent

from orchestration.airflow.engine.extract_assets import (
    ExtractAssetCommand,
    execute_extract_command,
    manifest_event_extra,
)


class TiDBCloudExportTrigger(BaseTrigger):
    def __init__(
        self,
        *,
        command: dict[str, object],
        export_id: str,
        poll_interval_seconds: int,
        timeout_seconds: int,
    ) -> None:
        super().__init__()
        self.command = command
        self.export_id = export_id
        self.poll_interval_seconds = poll_interval_seconds
        self.timeout_seconds = timeout_seconds

    def serialize(self) -> tuple[str, dict[str, object]]:
        return (
            "orchestration.airflow.engine.tidb_extract_operator.TiDBCloudExportTrigger",
            {
                "command": self.command,
                "export_id": self.export_id,
                "poll_interval_seconds": self.poll_interval_seconds,
                "timeout_seconds": self.timeout_seconds,
            },
        )

    async def run(self) -> AsyncIterator[TriggerEvent]:
        command = ExtractAssetCommand.from_trigger_kwargs(self.command)
        started_at = datetime.now(UTC)
        while True:
            status = await asyncio.to_thread(
                execute_extract_command,
                command,
                "export-status",
                export_id=self.export_id,
            )
            export = status.get("export", {})
            state = str(status.get("state") or export.get("state") or "")
            if state == "SUCCEEDED":
                yield TriggerEvent({"status": "success", "export_id": self.export_id, "state": state})
                return
            if state in {"FAILED", "CANCELED"}:
                yield TriggerEvent({"status": "failed", "export_id": self.export_id, "state": state})
                return
            elapsed = (datetime.now(UTC) - started_at).total_seconds()
            if elapsed >= self.timeout_seconds:
                yield TriggerEvent({"status": "timeout", "export_id": self.export_id, "state": state})
                return
            await asyncio.sleep(self.poll_interval_seconds)


class TiDBCloudExportManifestOperator(BaseOperator):
    def __init__(
        self,
        *,
        command: ExtractAssetCommand,
        manifest_asset: Any,
        poll_interval_seconds: int,
        timeout_seconds: int,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.command = command.to_trigger_kwargs()
        self.manifest_asset = manifest_asset
        self.poll_interval_seconds = poll_interval_seconds
        self.timeout_seconds = timeout_seconds

    def execute(self, context: dict[str, Any]) -> None:
        command = ExtractAssetCommand.from_trigger_kwargs(self.command)
        export = execute_extract_command(command, "start-export")
        export_id = export.get("export_id")
        if not export_id:
            raise RuntimeError("TiDB Cloud start-export did not return an export_id")
        self.defer(
            trigger=TiDBCloudExportTrigger(
                command=self.command,
                export_id=str(export_id),
                poll_interval_seconds=self.poll_interval_seconds,
                timeout_seconds=self.timeout_seconds,
            ),
            method_name="execute_complete",
            kwargs={
                "run_date": export["run_date"],
                "run_id": export["run_id"],
                "export_id": str(export_id),
                "start_payload": export,
            },
        )

    def execute_complete(
        self,
        context: dict[str, Any],
        event: dict[str, Any],
        *,
        run_date: str,
        run_id: str,
        export_id: str,
        start_payload: dict[str, Any],
    ) -> dict[str, object]:
        if event["status"] != "success":
            raise RuntimeError(
                f"TiDB Cloud export {export_id} ended with status={event['status']} state={event.get('state')}"
            )
        command = ExtractAssetCommand.from_trigger_kwargs(self.command)
        manifest = execute_extract_command(
            command,
            "write-manifest",
            run_date=run_date,
            run_id=run_id,
            export_id=export_id,
        )
        extra = manifest_event_extra(manifest_payload=manifest, fallback_payload=start_payload)
        context["outlet_events"][self.manifest_asset].extra = extra
        return {"manifest": extra, "export": {"export_id": export_id, "state": event.get("state")}}