import json
import logging
import os
import signal
import subprocess
import sys
import threading
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager, suppress
from types import MappingProxyType
from typing import Any, Final, assert_never

from mypy_boto3_sqs.type_defs import MessageTypeDef
from pydantic import ValidationError

from src.worker import task_protection
from src.worker.clients import sqs, step_functions
from src.worker.clients.step_functions import StaleTaskToken
from src.worker.exceptions import (
    AbortedStaleTask,
    MalformedMessage,
    MalformedSubprocessOutput,
    SubprocessCrashError,
    UnknownSubprocessError,
)
from src.worker.request import WorkerRequest
from src.worker.subprocess import WORKER_REQUEST, ErrorOutput
from src.worker.task_protection import TaskProtectionResult

_WORKER_EXCEPTIONS: Final = frozenset(
    {
        AbortedStaleTask,
        MalformedMessage,
        SubprocessCrashError,
        UnknownSubprocessError,
        MalformedSubprocessOutput,
    }
)

_logger = logging.getLogger(__name__)


class Worker[T_Request: WorkerRequest]:
    def __init__(
        self,
        *,
        request_schema: type[T_Request],
        select_subprocess_module: Callable[[T_Request], str],
        sqs_queue_url: str,
        heartbeat_interval_seconds: float = 15.0,
        message_processing_timeout_minutes: int = 30,
        exception_levels: Mapping[type[Exception], int] = MappingProxyType({}),
    ) -> None:
        self._request_schema = request_schema
        self._select_subprocess_module = select_subprocess_module
        self._sqs_queue_url = sqs_queue_url
        self._exception_levels = exception_levels
        # Worker-owned exceptions are excluded — the worker raises them
        # directly, never dispatched by name from subprocess output.
        self._subprocess_exception_for_name = {
            cls.__name__: cls for cls in exception_levels if cls not in _WORKER_EXCEPTIONS
        }
        self._heartbeat_interval_seconds = heartbeat_interval_seconds
        self._message_processing_timeout_minutes = message_processing_timeout_minutes

    def run(self) -> None:
        shutdown = threading.Event()
        # ECS defers SIGTERM on active task protection, so it normally arrives
        # only between messages (forced stops can bypass).
        signal.signal(signal.SIGTERM, lambda *_: shutdown.set())

        self.process_messages(shutdown)

    def process_messages(self, shutdown: threading.Event) -> None:
        _logger.info("worker starting")

        while not shutdown.is_set():
            message = sqs.receive_message(queue_url=self._sqs_queue_url)
            if message is None:
                continue
            if not self._process_message(message):
                break

        _logger.info("worker exiting")

    def _process_message(self, raw_message: MessageTypeDef) -> bool:
        result = task_protection.enable(expires_in_minutes=self._message_processing_timeout_minutes)

        # Returning False stops the polling loop so ECS can drain/replace the
        # task; the released or undeleted message gets picked up by another
        # worker on its next visibility cycle.
        match result:
            case (
                TaskProtectionResult.DEPLOYMENT_BLOCKED
                | TaskProtectionResult.TASK_STOPPING_OR_STOPPED
                | TaskProtectionResult.FAILED
            ):
                _logger.info("task protection unavailable (result=%s), releasing message", result)
                sqs.release_message(queue_url=self._sqs_queue_url, message=raw_message)
                return False
            case TaskProtectionResult.ENABLED:
                pass
            case _:
                assert_never(result)

        try:
            sqs.delete_message(queue_url=self._sqs_queue_url, message=raw_message)
            self._handle_message(raw_message["Body"])
        finally:
            disabled = task_protection.disable()
            if not disabled:
                _logger.error("failed to disable task protection")

        return disabled

    def _handle_message(self, body: str) -> None:
        try:
            request = self._parse_request(body)
        except MalformedMessage as exc:
            _logger.exception("message schema invalid. body=%s", body)
            token = self._extract_task_token(body)
            if token is not None:
                self._send_task_failure(token, exc)
            return

        try:
            step_functions.send_task_heartbeat(task_token=request.task_token)
        except StaleTaskToken:
            _logger.info("task token already stale at pickup; skipping")
            return
        except Exception as exc:
            _logger.warning("initial heartbeat failed; proceeding anyway: %s", exc, exc_info=exc)

        try:
            result = self._run_subprocess(request, body)
            self._send_task_success(request.task_token, result)
        except AbortedStaleTask:
            _logger.info("task token went stale; subprocess aborted")
        except Exception as exc:
            level = self._exception_levels.get(type(exc), logging.ERROR)
            self._send_task_failure(request.task_token, exc, level=level)

    def _parse_request(self, body: str) -> T_Request:
        try:
            return self._request_schema.model_validate_json(body)
        except ValidationError as exc:
            raise MalformedMessage(str(exc)) from exc

    @staticmethod
    def _extract_task_token(body: str) -> str | None:
        try:
            parsed = json.loads(body)
        except json.JSONDecodeError:
            return None
        if not isinstance(parsed, dict):
            return None
        token = parsed.get("task_token")
        if isinstance(token, str) and token:
            return token
        return None

    def _run_subprocess(self, request: T_Request, body: str) -> dict[str, Any]:
        process = self._spawn_subprocess(request, body)
        aborted_stale_task = threading.Event()
        with self._heartbeat(request.task_token, process, aborted_stale_task):
            stdout, stderr = process.communicate()
        return self._parse_subprocess_result(process, stdout, stderr, aborted_stale_task=aborted_stale_task.is_set())

    def _spawn_subprocess(self, request: T_Request, body: str) -> subprocess.Popen[str]:
        # Run the subprocess in a child process group so the heartbeat thread
        # can terminate the entire subprocess tree (Python orchestrator + any
        # tools it spawned) with a single killpg when the SFN task token goes
        # stale.
        env = os.environ.copy()
        env[WORKER_REQUEST] = body
        return subprocess.Popen(
            [sys.executable, "-m", self._select_subprocess_module(request)],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            start_new_session=True,
            env=env,
        )

    @contextmanager
    def _heartbeat(
        self,
        task_token: str,
        process: subprocess.Popen[str],
        aborted_stale_task: threading.Event,
    ) -> Iterator[None]:
        stop = threading.Event()
        thread = threading.Thread(
            target=self._run_heartbeat,
            args=(task_token, stop, process, aborted_stale_task, self._heartbeat_interval_seconds),
            daemon=True,
        )
        thread.start()
        try:
            yield
        finally:
            stop.set()
            thread.join()

    @staticmethod
    def _run_heartbeat(
        task_token: str,
        stop: threading.Event,
        process: subprocess.Popen[str],
        aborted_stale_task: threading.Event,
        interval_seconds: float,
    ) -> None:
        try:
            while not stop.wait(interval_seconds):
                try:
                    step_functions.send_task_heartbeat(task_token=task_token)
                except StaleTaskToken:
                    _logger.info("task token stale; terminating subprocess")
                    aborted_stale_task.set()
                    Worker._terminate_process_group(process)
                    return
                except Exception as exc:
                    _logger.warning(
                        "step_functions send_task_heartbeat failed task_token=%s: %s",
                        task_token,
                        exc,
                        exc_info=exc,
                    )
        except Exception:
            _logger.exception("heartbeat thread dying unexpectedly task_token=%s", task_token)

    @staticmethod
    def _terminate_process_group(process: subprocess.Popen[str]) -> None:
        if process.poll() is None:
            with suppress(ProcessLookupError):
                os.killpg(os.getpgid(process.pid), signal.SIGKILL)

    def _parse_subprocess_result(
        self,
        process: subprocess.Popen[str],
        stdout: str,
        stderr: str,
        *,
        aborted_stale_task: bool,
    ) -> dict[str, Any]:
        if stderr:
            stderr_level = logging.INFO if process.returncode == 0 else logging.WARNING
            _logger.log(stderr_level, "subprocess stderr: %s", stderr.strip())

        if process.returncode < 0:
            if aborted_stale_task:
                raise AbortedStaleTask("task token went stale")
            raise SubprocessCrashError(f"subprocess terminated by signal {-process.returncode}")

        if process.returncode == 0:
            try:
                parsed = json.loads(stdout)
            except json.JSONDecodeError as exc:
                raise MalformedSubprocessOutput(f"subprocess output was not valid JSON: {stdout!r}") from exc
            if not isinstance(parsed, dict):
                raise MalformedSubprocessOutput(f"subprocess output was not a JSON object: {stdout!r}")
            return parsed

        error_output = self._parse_subprocess_error_output(stdout, stderr, process.returncode)

        exc_class = self._subprocess_exception_for_name.get(error_output.error)
        if exc_class is None:
            raise UnknownSubprocessError(f"unknown error={error_output.error!r}: {error_output.message}")
        raise exc_class(error_output.message)

    @staticmethod
    def _parse_subprocess_error_output(stdout: str, stderr: str, returncode: int) -> ErrorOutput:
        try:
            return ErrorOutput.model_validate_json(stdout.strip())
        except ValidationError as e:
            raise MalformedSubprocessOutput(
                f"subprocess emitted unparseable error envelope (exit={returncode}): stdout={stdout!r} stderr={stderr.strip()!r}"
            ) from e

    def _send_task_success(self, task_token: str, result: dict[str, Any]) -> None:
        try:
            step_functions.send_task_success(task_token=task_token, output=result)
        except StaleTaskToken:
            _logger.info("task token stale before success could be reported")
        except Exception:
            _logger.exception("step_functions send_task_success failed; lost result: %s", json.dumps(result))

    def _send_task_failure(self, task_token: str, exc: Exception, *, level: int = logging.ERROR) -> None:
        _logger.log(level, "%s: %s", type(exc).__name__, exc, exc_info=exc)
        try:
            step_functions.send_task_failure(task_token=task_token, error=type(exc).__name__, cause=str(exc))
        except StaleTaskToken:
            _logger.info("task token stale before failure could be reported")
        except Exception:
            _logger.exception("step_functions send_task_failure failed")