"""Casty worker service for remote nodes.

Runs as a long-lived process on each node in the cluster.
Starts a ClusteredActorSystem and spawns discoverable worker actors
for direct task execution via ClusterClient.

Architecture:
- Each node spawns a local worker actor wrapped in Behaviors.discoverable().
- Workers register with WORKER_KEY for service discovery.
- The executor discovers workers via ClusterClient.lookup(WORKER_KEY).
- All nodes are symmetric (no head-only HTTP server).
"""

from __future__ import annotations

import argparse
import asyncio
import os
import traceback
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import Any

from casty import (
    ActorContext, ActorRef, Behavior, Behaviors, ServiceKey,
    CastyConfig, FailureDetectorConfig, HeartbeatConfig,
    ClusteredActorSystem
)
from loguru import logger


@dataclass(frozen=True, slots=True)
class TaskSucceeded:
    result: Any
    node_id: int


@dataclass(frozen=True, slots=True)
class TaskFailed:
    error: str
    traceback: str
    node_id: int


type TaskResult = TaskSucceeded | TaskFailed


@dataclass(frozen=True, slots=True)
class ExecuteTask:
    fn_bytes: bytes
    reply_to: ActorRef[TaskResult]


@dataclass(frozen=True, slots=True)
class _TaskDone:
    result: TaskResult
    reply_to: ActorRef[TaskResult]


@dataclass(frozen=True, slots=True)
class _TaskErrored:
    error: Exception
    reply_to: ActorRef[TaskResult]


type WorkerMsg = ExecuteTask | _TaskDone | _TaskErrored

WORKER_KEY: ServiceKey[WorkerMsg] = ServiceKey("skyward-worker")


def worker_behavior(node_id: int, concurrency: int = 1) -> Behavior[WorkerMsg]:
    log = logger.bind(component="worker", node_id=node_id)
    sem = asyncio.Semaphore(concurrency)

    async def _execute(fn_bytes: bytes) -> TaskResult:
        async with sem:
            def _run() -> TaskResult:
                from skyward.infra.serialization import deserialize

                try:
                    payload = deserialize(fn_bytes)
                    fn = payload["fn"]
                    args = payload.get("args", ())
                    kwargs = payload.get("kwargs", {})
                    fn_name = getattr(fn, "__name__", str(fn))

                    print(f"[worker-{node_id}] Executing {fn_name}...", flush=True)
                    result = fn(*args, **kwargs)
                    print(f"[worker-{node_id}] {fn_name} completed", flush=True)
                    return TaskSucceeded(result=result, node_id=node_id)
                except Exception as e:
                    fn_name = "unknown"
                    print(f"[worker-{node_id}] {fn_name} failed: {e}", flush=True)
                    return TaskFailed(
                        error=str(e),
                        traceback=traceback.format_exc(),
                        node_id=node_id,
                    )

            return await asyncio.to_thread(_run)

    async def receive(ctx: ActorContext[WorkerMsg], msg: WorkerMsg) -> Behavior[WorkerMsg]:
        match msg:
            case ExecuteTask(fn_bytes=fn_bytes, reply_to=reply_to):
                log.debug("ExecuteTask received, payload={size} bytes", size=len(fn_bytes))
                ctx.pipe_to_self(
                    coro=_execute(fn_bytes),
                    mapper=lambda result: _TaskDone(result=result, reply_to=reply_to),
                    on_failure=lambda e: _TaskErrored(error=e, reply_to=reply_to),
                )
                return Behaviors.same()
            case _TaskDone(result=result, reply_to=reply_to):
                log.debug("Task completed successfully")
                reply_to.tell(result)
                return Behaviors.same()
            case _TaskErrored(error=error, reply_to=reply_to):
                log.debug("Task errored: {error}", error=error)
                reply_to.tell(TaskFailed(
                    error=str(error),
                    traceback=traceback.format_exc(),
                    node_id=node_id,
                ))
                return Behaviors.same()

    return Behaviors.receive(receive)


async def main(
    node_id: int,
    port: int,
    seeds: list[tuple[str, int]] | None,
    num_nodes: int = 1,
    host: str = "0.0.0.0",
    workers_per_node: int = 1,
) -> None:
    config = CastyConfig(
        heartbeat=HeartbeatConfig(interval=2.0, availability_check_interval=5.0),
        failure_detector=FailureDetectorConfig(
            threshold=16.0,
            acceptable_heartbeat_pause_ms=10_000.0,
        ),
        suppress_dead_letters_on_shutdown=True
    )

    log = logger.bind(component="worker", node_id=node_id)
    quorum = num_nodes if num_nodes > 1 else None
    log.debug("Starting casty worker, quorum={quorum} port={port}", quorum=quorum, port=port)
    print(f"Casty worker {node_id} starting (quorum={quorum})...", flush=True)

    os.environ["SKYWARD_NODE_ID"] = str(node_id)

    async with ClusteredActorSystem(
        name="skyward",
        host=host,
        port=port,
        node_id=f"node-{node_id}",
        seed_nodes=tuple(seeds) if seeds else None,
        bind_host="0.0.0.0",
        config=config,
        required_quorum=quorum,
    ) as system:
        from skyward.distributed import _set_active_registry
        from skyward.distributed.proxies import set_system_loop
        from skyward.distributed.registry import DistributedRegistry

        loop = asyncio.get_running_loop()
        pool_size = max((os.cpu_count() or 1) + 4, workers_per_node)
        loop.set_default_executor(ThreadPoolExecutor(max_workers=pool_size))
        set_system_loop(loop)
        _set_active_registry(DistributedRegistry(system, loop=loop))

        system.spawn(
            Behaviors.discoverable(
                worker_behavior(node_id, concurrency=workers_per_node),
                key=WORKER_KEY,
            ),
            "worker",
        )
        print(
            f"Casty worker {node_id} ready "
            f"(cluster={num_nodes} nodes, concurrency={workers_per_node})",
            flush=True,
        )

        await asyncio.Event().wait()


def _parse_seeds(seeds_str: str | None) -> list[tuple[str, int]] | None:
    if not seeds_str:
        return None
    return [
        (host, int(port_str))
        for addr in seeds_str.split(",")
        for host, port_str in [addr.rsplit(":", 1)]
    ]


def cli() -> None:
    parser = argparse.ArgumentParser(description="Casty worker service")
    parser.add_argument("--node-id", type=int, required=True)
    parser.add_argument("--port", type=int, default=25520)
    parser.add_argument("--num-nodes", type=int, default=1)
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument(
        "--seeds", type=str, default=None,
        help="Comma-separated seed addresses (host:port)",
    )
    parser.add_argument("--workers-per-node", type=int, default=1)
    args = parser.parse_args()

    seeds = _parse_seeds(args.seeds)
    asyncio.run(main(
        args.node_id, args.port, seeds, args.num_nodes, args.host,
        workers_per_node=args.workers_per_node,
    ))


if __name__ == "__main__":
    # When running as `python -m skyward.infra.worker`, the import chain
    # (skyward → api.pool → infra.executor → infra.worker) loads this module
    # as 'skyward.infra.worker' first, then Python re-executes it as __main__.
    # This creates duplicate class objects: __main__.TaskSucceeded vs
    # skyward.infra.worker.TaskSucceeded. Pickle stores the __main__ version,
    # which the client can't resolve (its __main__ is a different module).
    # Fix: delegate to the properly-imported module so all class references
    # point to 'skyward.infra.worker', not '__main__'.
    import sys

    _proper = sys.modules.get("skyward.infra.worker")
    if _proper is not None:
        _proper.cli()
    else:
        cli()
