"""Ambient worker daemon — pulls jobs when idle, processes locally."""

from __future__ import annotations

import platform
import signal
import threading
import time
from concurrent.futures import Future, ThreadPoolExecutor
from contextlib import contextmanager
from copy import deepcopy
from dataclasses import dataclass
from typing import Any

import structlog

from hive_compute.model_ids import canonical_model_identifier
from hive_compute.models import WorkerInfo, WorkerStatus
from hive_compute.worker.api_client import CoordinatorClient
from hive_compute.worker.idle_detector import get_hardware_info, should_run
from hive_compute.worker.keychain import resolve_worker_token

logger = structlog.get_logger()


@dataclass
class _StagedWork:
    job: Any
    prepared: Any | None = None
    stage_source: str = "inline"

    def cleanup(self) -> None:
        if self.prepared is not None and hasattr(self.prepared, "cleanup"):
            self.prepared.cleanup()


class WorkerDaemon:
    """Main daemon loop: poll for jobs, process when idle, push results."""

    def __init__(
        self,
        config_path: str = "~/.ambient/config.yaml",
        runtime_overrides: dict[str, Any] | None = None,
    ):
        self.config_path = config_path
        self.runtime_overrides = runtime_overrides or {}
        self.config = self._load_config(config_path)
        self.client = self._build_client(self.config)
        self.handlers: dict[str, Any] = {}
        self._staging_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="hive-stage")
        self._staging_future: Future[_StagedWork | None] | None = None
        self._staging_cancel: threading.Event | None = None
        self._staged_work: _StagedWork | None = None
        self._shutdown = threading.Event()
        self._register_signal_handlers()
        self._sync_handlers(self.config)
        self._prewarm_handlers(reason="startup")

    def _load_config(self, path: str) -> dict:
        from hive_compute.worker.config_model import load_worker_config_raw

        config = load_worker_config_raw(path)
        config["worker_token"] = resolve_worker_token(config)
        return self._apply_runtime_overrides(config)

    def _apply_runtime_overrides(self, config: dict[str, Any]) -> dict[str, Any]:
        merged = deepcopy(config)
        for key, value in self.runtime_overrides.items():
            current = merged.get(key)
            if isinstance(current, dict) and isinstance(value, dict):
                current.update(value)
                merged[key] = current
            else:
                merged[key] = value
        return merged

    @staticmethod
    def _build_client(config: dict[str, Any]) -> CoordinatorClient:
        return CoordinatorClient(
            base_url=config["coordinator_url"],
            worker_token=config["worker_token"],
            worker_id=config.get("worker_id", platform.node()),
        )

    def _register_signal_handlers(self):
        signal.signal(signal.SIGTERM, self._handle_signal)
        signal.signal(signal.SIGINT, self._handle_signal)

    def _handle_signal(self, signum, frame):
        logger.info("Shutdown signal received", signal=signum)
        self._shutdown.set()

    def _sync_handlers(self, config: dict[str, Any]) -> bool:
        """Load or refresh configured job handlers."""
        handler_config = config.get("handlers", {})
        changed = False

        if "transcribe" in handler_config:
            from hive_compute.handlers.whisper_handler import WhisperHandler

            transcribe_config = handler_config["transcribe"]
            existing = self.handlers.get("transcribe")
            if isinstance(existing, WhisperHandler):
                if existing.update_config(transcribe_config):
                    changed = True
                    logger.info(
                        "Reloaded handler config",
                        type="transcribe",
                        model=existing.model,
                        backend=existing.backend,
                    )
            else:
                self.handlers["transcribe"] = WhisperHandler(transcribe_config)
                changed = True
                logger.info("Loaded handler", type="transcribe")
        elif "transcribe" in self.handlers:
            del self.handlers["transcribe"]
            changed = True
        return changed

    def _reload_runtime_config_if_changed(self):
        try:
            next_config = self._load_config(self.config_path)
        except Exception as exc:
            logger.warning("Failed to reload worker config", error=str(exc))
            return

        if next_config == self.config:
            return

        previous_identity = (
            self.config.get("coordinator_url"),
            self.config.get("worker_token"),
            self.config.get("worker_id"),
        )
        next_identity = (
            next_config.get("coordinator_url"),
            next_config.get("worker_token"),
            next_config.get("worker_id"),
        )

        self.config = next_config
        self._cancel_prefetch("worker config changed")
        handlers_changed = self._sync_handlers(next_config)
        if previous_identity != next_identity:
            self.client = self._build_client(next_config)
            handlers_changed = True

        logger.info("Worker config reloaded", worker_id=self.config.get("worker_id", platform.node()))
        if handlers_changed:
            self._prewarm_handlers(reason="config_reload")
            self._announce_worker_state()

    def _announce_worker_state(self):
        """Update worker info with the coordinator."""
        hw = get_hardware_info()
        self.client.update_worker(
            WorkerInfo(
                worker_id=self.config.get("worker_id", platform.node()),
                status=WorkerStatus.IDLE,
                available_handlers=list(self.handlers.keys()),
                available_models=self._transcribe_supported_models(),
                model_capabilities=self._transcribe_model_capabilities(),
                chip=hw.get("chip"),
                memory_gb=hw.get("memory_gb"),
            ).model_dump()
        )

    def _transcribe_supported_models(self) -> list[str]:
        transcribe_config = self.config.get("handlers", {}).get("transcribe", {})
        configured_models = list(transcribe_config.get("supported_models") or [])

        primary_model = transcribe_config.get("model")
        if primary_model:
            configured_models.append(
                canonical_model_identifier(
                    primary_model,
                    transcribe_config.get("quantization"),
                )
            )

        supported_models: list[str] = []
        seen: set[str] = set()
        for model in configured_models:
            normalized = str(model).strip()
            if not normalized:
                continue
            identifier = canonical_model_identifier(normalized, default_quantization="none")
            if identifier in seen:
                continue
            seen.add(identifier)
            supported_models.append(identifier)
        return supported_models

    def _transcribe_model_capabilities(self) -> list[dict[str, Any]]:
        handler = self.handlers.get("transcribe")
        if handler is None or not hasattr(handler, "model_capabilities"):
            return []
        try:
            return handler.model_capabilities(self._transcribe_supported_models())
        except Exception as exc:
            logger.warning("Failed to collect model capabilities", error=str(exc))
            return []

    def _prewarm_handlers(self, *, reason: str) -> None:
        handler = self.handlers.get("transcribe")
        if handler is None or not getattr(handler, "prewarm_model_on_startup", False):
            return
        if not hasattr(handler, "prewarm"):
            return

        result = handler.prewarm(reason=reason)
        status = result.get("status", "unknown")
        log_data = {
            "model": result.get("model"),
            "backend": result.get("backend"),
            "repo": result.get("repo"),
            "reason": result.get("reason"),
            "elapsed_seconds": result.get("elapsed_seconds"),
            "error": result.get("error"),
        }
        if status in {"warmed", "already_warm"}:
            logger.info("Worker model prewarm finished", status=status, **log_data)
        elif status == "skipped":
            logger.info("Worker model prewarm skipped", status=status, **log_data)
        else:
            logger.warning("Worker model prewarm failed", status=status, **log_data)

    def _take_staged_work(self) -> _StagedWork | None:
        if self._staged_work is not None:
            staged = self._staged_work
            self._staged_work = None
            return staged

        if self._staging_future is None or not self._staging_future.done():
            return None

        try:
            staged = self._staging_future.result()
        except Exception as exc:
            logger.warning("Background staging failed", error=str(exc))
            staged = None
        finally:
            self._staging_future = None
            self._staging_cancel = None
        return staged

    def _prepare_work_inline(self, job) -> _StagedWork:
        handler = self.handlers.get(job.job_type)
        if handler is None or not hasattr(handler, "prepare_job"):
            return _StagedWork(job=job, prepared=None, stage_source="inline")

        started_at = time.perf_counter()
        prepared = handler.prepare_job(job)
        logger.info(
            "Prepared job",
            job_id=job.id,
            shard_id=job.shard_id,
            stage_source="inline",
            stage_elapsed_seconds=round(time.perf_counter() - started_at, 3),
            **getattr(prepared, "stage_timings", {}),
        )
        return _StagedWork(job=job, prepared=prepared, stage_source="inline")

    def _start_prefetch(self) -> None:
        if self._shutdown.is_set() or self._staged_work is not None or self._staging_future is not None:
            return

        self._staging_cancel = threading.Event()
        self._staging_future = self._staging_executor.submit(self._prefetch_next_work, self._staging_cancel)

    def _prefetch_next_work(self, cancel_event: threading.Event) -> _StagedWork | None:
        if cancel_event.is_set():
            return None

        job = self.client.pull_job(list(self.handlers.keys()))
        if not job:
            return None

        if cancel_event.is_set():
            self._release_job(job, reason="prefetch_canceled_before_stage")
            return None

        try:
            staged = self._prepare_work_inline(job)
        except Exception as exc:
            logger.warning("Prefetch staging failed", job_id=job.id, shard_id=job.shard_id, error=str(exc))
            self.client.report_failure(job.id, str(exc), shard_id=job.shard_id)
            return None

        if cancel_event.is_set():
            staged.cleanup()
            self._release_job(job, reason="prefetch_canceled_after_stage")
            return None

        staged.stage_source = "prefetched"
        return staged

    def _cancel_prefetch(self, reason: str) -> None:
        if self._staging_cancel is not None:
            self._staging_cancel.set()
        if self._staging_future is not None and self._staging_future.done():
            try:
                staged = self._staging_future.result()
            except Exception:
                staged = None
            self._staging_future = None
            self._staging_cancel = None
            if staged is not None:
                staged.cleanup()
                self._release_job(staged.job, reason=reason)
        if self._staged_work is not None:
            staged = self._staged_work
            self._staged_work = None
            staged.cleanup()
            self._release_job(staged.job, reason=reason)

    def _release_job(self, job, *, reason: str) -> None:
        if not hasattr(self.client, "release_job"):
            return
        released = self.client.release_job(job.id, shard_id=job.shard_id, reason=reason)
        if not released:
            logger.warning("Failed to release staged job", job_id=job.id, shard_id=job.shard_id, reason=reason)

    def _is_paused(self) -> bool:
        return False

    def _pause_wait(self) -> None:
        self._shutdown.wait(2)

    def _on_worker_started(self) -> None:
        """Hook for subclasses that need to react after startup announcement."""

    def _on_waiting_for_idle(self) -> None:
        """Hook for subclasses while the worker is ineligible to run."""

    def _on_waiting_for_job(self) -> None:
        """Hook for subclasses while the worker is polling with no work."""

    def _on_job_started(self, job) -> None:
        """Hook for subclasses when a job begins processing."""

    def _on_job_completed(self, job, result: dict[str, Any]) -> None:
        """Hook for subclasses after a successful result submission."""

    def _on_job_failed(self, job, error: Exception) -> None:
        """Hook for subclasses after a job failure is reported."""

    def _on_job_finished(self, job) -> None:
        """Hook for subclasses after a job leaves the working state."""

    def _on_worker_stopped(self) -> None:
        """Hook for subclasses when the daemon exits."""

    def run(self):
        """Main loop."""
        logger.info(
            "Worker daemon starting",
            worker_id=self.config.get("worker_id", platform.node()),
            handlers=list(self.handlers.keys()),
        )

        self._announce_worker_state()
        self._on_worker_started()

        try:
            while not self._shutdown.is_set():
                if self._is_paused():
                    self._cancel_prefetch("worker_paused")
                    self._pause_wait()
                    continue

                self._reload_runtime_config_if_changed()
                idle_config = self.config.get("idle_detection", {})
                poll_config = self.config.get("polling", {})
                idle_interval = poll_config.get("idle_interval_seconds", 60)
                poll_interval = poll_config.get("interval_seconds", 30)
                busy_interval = poll_config.get("busy_interval_seconds", 5)

                if not should_run(
                    min_idle_minutes=idle_config.get("min_idle_minutes", 5),
                    require_ac=idle_config.get("require_ac_power", True),
                    check_thermal=idle_config.get("check_thermal", True),
                ):
                    self._cancel_prefetch("worker_became_ineligible")
                    self._on_waiting_for_idle()
                    self._shutdown.wait(idle_interval)
                    continue

                staged_work = self._take_staged_work()
                if staged_work is None:
                    job = self.client.pull_job(list(self.handlers.keys()))
                    if not job:
                        self._on_waiting_for_job()
                        self._shutdown.wait(poll_interval)
                        continue
                    try:
                        staged_work = self._prepare_work_inline(job)
                    except Exception as exc:
                        logger.error("Failed to prepare job", job_id=job.id, shard_id=job.shard_id, error=str(exc))
                        self.client.report_failure(job.id, str(exc), shard_id=job.shard_id)
                        continue

                job = staged_work.job
                handler = self.handlers.get(job.job_type)
                if not handler:
                    logger.error("No handler for job type", job_type=job.job_type)
                    staged_work.cleanup()
                    self.client.report_failure(job.id, f"No handler for {job.job_type}", shard_id=job.shard_id)
                    continue

                logger.info(
                    "Processing job",
                    job_id=job.id,
                    job_type=job.job_type,
                    shard_id=job.shard_id,
                    stage_source=staged_work.stage_source,
                )
                self._start_prefetch()
                self._on_job_started(job)
                try:
                    with self._heartbeat_context(job.id, job.shard_id):
                        if staged_work.prepared is not None and hasattr(handler, "run_prepared"):
                            result = handler.run_prepared(staged_work.prepared)
                        else:
                            result = handler.run(job)
                    submit_started_at = time.perf_counter()
                    self.client.push_result(job.id, result, shard_id=job.shard_id)
                    submit_elapsed = time.perf_counter() - submit_started_at
                    pipeline_timings = dict(result.get("pipeline_timings") or {})
                    pipeline_timings["submit_seconds"] = round(submit_elapsed, 3)
                    logger.info(
                        "Job completed",
                        job_id=job.id,
                        shard_id=job.shard_id,
                        stage_source=staged_work.stage_source,
                        **pipeline_timings,
                    )
                    self._on_job_completed(job, result)
                except Exception as exc:
                    logger.error("Job failed", job_id=job.id, shard_id=job.shard_id, error=str(exc))
                    self.client.report_failure(job.id, str(exc), shard_id=job.shard_id)
                    self._on_job_failed(job, exc)
                finally:
                    staged_work.cleanup()
                    self._on_job_finished(job)

                self._shutdown.wait(busy_interval)
        finally:
            self._cancel_prefetch("worker_shutdown")
            self._staging_executor.shutdown(wait=False, cancel_futures=True)
            logger.info("Worker daemon stopped")
            self._on_worker_stopped()

    @contextmanager
    def _heartbeat_context(self, job_id: str, shard_id: str | None = None):
        """Send periodic heartbeats while processing a job or shard."""
        stop = threading.Event()

        def beat():
            while not stop.is_set():
                self.client.heartbeat(job_id, shard_id=shard_id)
                stop.wait(60)

        t = threading.Thread(target=beat, daemon=True)
        t.start()
        try:
            yield
        finally:
            stop.set()
