from __future__ import annotations

from dataclasses import dataclass, field
import time
from typing import Any, Callable, Literal

from flownet.api.services import get_services

TaskKind = Literal["service", "source", "subscription", "publisher"]

_TERMINAL_STATUSES = frozenset({"stopped", "error", "completed", "failed", "cancelled"})
_POLL_INTERVAL_SECONDS = 0.01
_KIND_ALIASES: dict[str, TaskKind] = {
    "service": "service",
    "services": "service",
    "source": "source",
    "sources": "source",
    "subscription": "subscription",
    "subscriptions": "subscription",
    "publisher": "publisher",
    "publish": "publisher",
}


@dataclass
class TaskHandle:
    kind: str
    raw_id: str
    _status_fn: Callable[[], str]
    _cancel_fn: Callable[[], bool]
    _wait_fn: Callable[[float | None], bool] | None = None
    meta: dict[str, Any] = field(default_factory=dict)

    def id(self) -> str:
        return _compose_task_id(self.kind, self.raw_id)

    def status(self) -> str:
        value = self._status_fn()
        return str(value or "unknown")

    def cancel(self) -> bool:
        return bool(self._cancel_fn())

    def wait(self, timeout: float | None = None) -> bool:
        if self._wait_fn is not None:
            return bool(self._wait_fn(timeout))
        return _wait_until(lambda: self.status() in _TERMINAL_STATUSES, timeout)

    def to_dict(self) -> dict[str, Any]:
        payload = {
            "id": self.id(),
            "kind": self.kind,
            "raw_id": self.raw_id,
            "status": self.status(),
        }
        if self.meta:
            payload["meta"] = dict(self.meta)
        return payload



def list(kind: str | None = None) -> list[TaskHandle]:
    normalized_kind = _normalize_kind(kind)
    if normalized_kind == "service":
        return _list_service_handles()
    if normalized_kind == "source":
        return _list_source_handles()
    if normalized_kind == "subscription":
        return _list_subscription_handles()
    if normalized_kind == "publisher":
        return _list_service_handles() + _list_source_handles()
    return _list_service_handles() + _list_source_handles() + _list_subscription_handles()



def get(task_id: str) -> TaskHandle | None:
    if not isinstance(task_id, str) or not task_id:
        raise ValueError("task_id must be a non-empty string.")

    kind_hint, raw_id = _split_task_id(task_id)
    if kind_hint is not None:
        return _get_by_kind(kind_hint, raw_id)

    for kind in ("subscription", "source", "service"):
        handle = _get_by_kind(kind, raw_id)
        if handle is not None:
            return handle
    return None



def cancel(task_id: str) -> bool:
    handle = get(task_id)
    if handle is None:
        return False
    return handle.cancel()



def _normalize_kind(kind: str | None) -> TaskKind | None:
    if kind is None:
        return None
    normalized = kind.strip().lower()
    if not normalized:
        return None
    value = _KIND_ALIASES.get(normalized)
    if value is None:
        raise ValueError(
            "kind must be one of: service, source, subscription, publisher"
        )
    return value



def _compose_task_id(kind: str, raw_id: str) -> str:
    return f"{kind}:{raw_id}"



def _split_task_id(task_id: str) -> tuple[str | None, str]:
    kind, sep, raw_id = task_id.partition(":")
    if not sep:
        return None, task_id
    kind_hint = _KIND_ALIASES.get(kind.strip().lower())
    if kind_hint is None:
        return None, task_id
    return kind_hint, raw_id



def _wait_until(predicate: Callable[[], bool], timeout: float | None) -> bool:
    if timeout is not None and timeout < 0:
        raise ValueError("timeout must be >= 0.")
    deadline = None if timeout is None else time.monotonic() + timeout
    while True:
        if predicate():
            return True
        if deadline is not None and time.monotonic() >= deadline:
            return False
        time.sleep(_POLL_INTERVAL_SECONDS)



def _list_service_handles() -> list[TaskHandle]:
    hub = get_services()
    entries = hub.list_services()
    handles: list[TaskHandle] = []
    for entry in entries:
        service_name = entry.get("name")
        if isinstance(service_name, str) and service_name:
            handles.append(_service_task_handle(service_name))
    return handles



def _list_source_handles() -> list[TaskHandle]:
    hub = get_services()
    entries = hub.list_sources()
    handles: list[TaskHandle] = []
    for entry in entries:
        source_id = entry.get("source_id")
        if isinstance(source_id, str) and source_id:
            handles.append(_source_task_handle(source_id))
    return handles



def _list_subscription_handles() -> list[TaskHandle]:
    hub = get_services()
    topic_registry = hub.require_topic_registry()
    list_subscriptions = getattr(topic_registry, "list_subscriptions", None)
    if not callable(list_subscriptions):
        return []

    handles: list[TaskHandle] = []
    for entry in list_subscriptions():
        subscription_id = entry.get("subscription_id")
        if not isinstance(subscription_id, str) or not subscription_id:
            continue
        metadata = {
            "topic_id": entry.get("topic_id"),
            "execution_address": entry.get("execution_address"),
            "tags": dict(entry.get("tags") or {}),
        }
        handles.append(_subscription_task_handle(subscription_id, metadata=metadata))
    return handles



def _get_by_kind(kind: str, raw_id: str) -> TaskHandle | None:
    if kind == "service":
        return _get_service_handle(raw_id)
    if kind == "source":
        return _get_source_handle(raw_id)
    if kind == "subscription":
        return _get_subscription_handle(raw_id)
    if kind == "publisher":
        # publisher is an aggregate group and does not map to a stable id space.
        return None
    return None



def _get_service_handle(name: str) -> TaskHandle | None:
    if not name:
        return None
    if get_services().require_service_registry().get_service(name) is None:
        return None
    return _service_task_handle(name)



def _get_source_handle(source_id: str) -> TaskHandle | None:
    if not source_id:
        return None
    if get_services().require_sources().registry.get_source(source_id) is None:
        return None
    return _source_task_handle(source_id)



def _get_subscription_handle(subscription_id: str) -> TaskHandle | None:
    if not subscription_id:
        return None
    topic_registry = get_services().require_topic_registry()
    get_subscription = getattr(topic_registry, "get_subscription", None)
    if not callable(get_subscription):
        return None
    subscription = get_subscription(subscription_id)
    if subscription is None:
        return None
    metadata = {
        "topic_id": getattr(subscription, "topic_id", None),
        "execution_address": getattr(subscription, "execution_address", None),
        "tags": dict(getattr(subscription, "tags", None) or {}),
    }
    return _subscription_task_handle(subscription_id, metadata=metadata)



def _service_task_handle(name: str) -> TaskHandle:
    def _status() -> str:
        handle = get_services().require_service_registry().get_service(name)
        if handle is None:
            return "stopped"
        return str(getattr(handle, "status", "unknown"))

    def _cancel() -> bool:
        return bool(get_services().stop_service(name))

    return TaskHandle(
        kind="service",
        raw_id=name,
        _status_fn=_status,
        _cancel_fn=_cancel,
        meta={"name": name},
    )



def _source_task_handle(source_id: str) -> TaskHandle:
    def _status() -> str:
        handle = get_services().require_sources().registry.get_source(source_id)
        if handle is None:
            return "stopped"
        return str(getattr(handle, "status", "unknown"))

    def _cancel() -> bool:
        return bool(get_services().stop_source(source_id))

    return TaskHandle(
        kind="source",
        raw_id=source_id,
        _status_fn=_status,
        _cancel_fn=_cancel,
        meta={"source_id": source_id},
    )



def _subscription_task_handle(
    subscription_id: str,
    *,
    metadata: dict[str, Any] | None = None,
) -> TaskHandle:
    def _status() -> str:
        topic_registry = get_services().require_topic_registry()
        get_subscription = getattr(topic_registry, "get_subscription", None)
        if not callable(get_subscription):
            return "unknown"
        return "running" if get_subscription(subscription_id) is not None else "stopped"

    def _cancel() -> bool:
        return bool(get_services().require_topic_registry().detach_task(subscription_id))

    return TaskHandle(
        kind="subscription",
        raw_id=subscription_id,
        _status_fn=_status,
        _cancel_fn=_cancel,
        meta=dict(metadata or {}),
    )


__all__ = [
    "TaskHandle",
    "TaskKind",
    "list",
    "get",
    "cancel",
]
