from __future__ import annotations

import asyncio
import base64
import hashlib
import sys
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any

import pytest
from pydantic import ValidationError

from cmvr_edge_ai.contracts import (
    ImageArtifact,
    ImageInput,
    InferenceRequest,
    InferenceStatus,
    JsonArtifact,
    ScalarOutput,
)
from cmvr_edge_ai.core import Envelope
from cmvr_edge_ai.gauge import (
    GAUGE_CATEGORY,
    GAUGE_MODEL_ID,
    AnalogGaugeReaderOperator,
)
from cmvr_edge_ai.workers import PROTOCOL_VERSION, FramedWorkerResponse


class FakeWorker:
    def __init__(
        self,
        responses: Sequence[FramedWorkerResponse | Exception],
        *,
        ready: Mapping[str, Any] | None = None,
    ) -> None:
        self.responses = list(responses)
        self.ready = dict(
            ready
            or {
                "protocol": PROTOCOL_VERSION,
                "type": "ready",
                "model_id": GAUGE_MODEL_ID,
            }
        )
        self.calls: list[tuple[Mapping[str, Any], tuple[bytes, ...], float]] = []
        self.is_alive = False
        self.last_error = ""
        self.stopped = False

    async def start(self) -> Mapping[str, Any]:
        self.is_alive = True
        return self.ready

    async def submit(
        self,
        header: Mapping[str, Any],
        blobs: Sequence[bytes] = (),
        *,
        timeout_s: float,
    ) -> FramedWorkerResponse:
        self.calls.append((dict(header), tuple(blobs), timeout_s))
        outcome = self.responses.pop(0)
        if isinstance(outcome, Exception):
            raise outcome
        return outcome

    async def stop(self) -> None:
        self.is_alive = False
        self.stopped = True


def _params(tmp_path: Path) -> dict[str, object]:
    model_dir = tmp_path / "models"
    model_dir.mkdir(exist_ok=True)
    for name in ("detection.pt", "keypoints.pt", "segmentation.pt"):
        (model_dir / name).write_bytes(b"model")
    return {
        "python_executable": sys.executable,
        "project_root": str(tmp_path),
        "detection_model_path": "models/detection.pt",
        "key_point_model_path": "models/keypoints.pt",
        "segmentation_model_path": "models/segmentation.pt",
        "startup_timeout_s": 12,
        "request_timeout_s": 34,
        "shutdown_timeout_s": 5,
        "max_header_bytes": 4096,
        "max_blob_bytes": 8192,
        "max_pixels": 1_234_567,
    }


def _request(
    *,
    request_id: str = "gauge-1",
    roles: tuple[str, ...] = ("original", "annotated", "diagnostics"),
) -> InferenceRequest:
    encoded = b"encoded-gauge-image"
    return InferenceRequest(
        request_id=request_id,
        category=GAUGE_CATEGORY,
        source_id="camera-front",
        inputs=(
            ImageInput(
                name="image",
                media_type="image/jpeg",
                data=base64.b64encode(encoded).decode("ascii"),
                width=640,
                height=480,
                sha256=hashlib.sha256(encoded).hexdigest(),
            ),
        ),
        requested_artifact_roles=roles,
    )


def _envelope(request: InferenceRequest) -> Envelope[InferenceRequest]:
    return Envelope(
        request,
        schema_name="InferenceRequest",
        trace_id=request.request_id or "fallback",
        source_id=request.source_id or "",
        attributes={
            "invocation_request_id": request.request_id,
            "invocation_category": request.category,
        },
    )


def _worker_response(
    *,
    request_id: str = "gauge-1",
    status: str = "succeeded",
    value: float | None = 12.5,
    with_annotated: bool = True,
    input_image: dict[str, int] | None = None,
) -> FramedWorkerResponse:
    annotated = b"annotated-jpeg"
    artifacts = (
        [
            {
                "blob_index": 0,
                "role": "annotated",
                "media_type": "image/jpeg",
                "width": 640,
                "height": 480,
                "sha256": hashlib.sha256(annotated).hexdigest(),
            }
        ]
        if with_annotated
        else []
    )
    error = (
        {
            "code": "GAUGE_READER_FAILED",
            "message": "gauge could not be read",
            "retryable": False,
            "stage": "ocr",
            "details": {"reason": "no labels"},
        }
        if status == "failed"
        else None
    )
    return FramedWorkerResponse(
        header={
            "protocol": PROTOCOL_VERSION,
            "type": "response",
            "request_id": request_id,
            "category": GAUGE_CATEGORY,
            "model_id": GAUGE_MODEL_ID,
            "status": status,
            "value": value,
            "unit": "bar" if value is not None else None,
            "inference_ms": 7.5,
            "input_image": (
                {"width": 640, "height": 480}
                if input_image is None and status != "failed"
                else input_image
            ),
            "diagnostics": {"ocr_labels": 4, "fit": {"residual": 0.02}},
            "warnings": [],
            "error": error,
            "artifacts": artifacts,
        },
        blobs=(annotated,) if with_annotated else (),
    )


def test_operator_maps_scalar_and_requested_artifacts(tmp_path: Path) -> None:
    async def exercise() -> None:
        fake = FakeWorker((_worker_response(),))
        constructed: dict[str, object] = {}

        def factory(command, **kwargs):  # type: ignore[no-untyped-def]
            constructed["command"] = tuple(command)
            constructed["kwargs"] = kwargs
            return fake

        operator = AnalogGaugeReaderOperator(
            "reader",
            _params(tmp_path),
            worker_factory=factory,
        )
        await operator.start()
        emission = await operator.process(_envelope(_request()))
        response = emission.envelope.payload

        assert response.status is InferenceStatus.SUCCEEDED
        assert response.request_id == "gauge-1"
        assert response.category == GAUGE_CATEGORY
        assert response.source_id == "camera-front"
        assert response.model is not None
        assert response.model.model_id == GAUGE_MODEL_ID
        assert response.outputs == (
            ScalarOutput(name="reading", value=12.5, unit="bar"),
        )
        assert [artifact.role for artifact in response.artifacts] == [
            "original",
            "annotated",
            "diagnostics",
        ]
        original = response.artifacts[0]
        annotated = response.artifacts[1]
        diagnostics = response.artifacts[2]
        assert isinstance(original, ImageArtifact)
        assert base64.b64decode(original.data) == b"encoded-gauge-image"
        assert isinstance(annotated, ImageArtifact)
        assert base64.b64decode(annotated.data) == b"annotated-jpeg"
        assert isinstance(diagnostics, JsonArtifact)
        assert diagnostics.data == {
            "ocr_labels": 4,
            "fit": {"residual": 0.02},
        }

        command = constructed["command"]
        assert isinstance(command, tuple)
        assert Path(command[2]).is_absolute()
        assert command[2].endswith("cmvr_edge_ai/gauge/analog_gauge_worker.py")
        max_pixels_index = command.index("--max-pixels")
        assert command[max_pixels_index + 1] == "1234567"
        sent_header, sent_blobs, timeout_s = fake.calls[0]
        assert sent_header["request_id"] == "gauge-1"
        assert sent_header["model_id"] == GAUGE_MODEL_ID
        assert sent_header["image"] == {
            "blob_index": 0,
            "name": "image",
            "role": "original",
            "media_type": "image/jpeg",
            "sha256": hashlib.sha256(b"encoded-gauge-image").hexdigest(),
            "width": 640,
            "height": 480,
        }
        assert sent_blobs == (b"encoded-gauge-image",)
        assert timeout_s == 34
        await operator.stop()
        assert fake.stopped is True

    asyncio.run(exercise())


def test_invalid_worker_response_fails_one_request_then_recovers(
    tmp_path: Path,
) -> None:
    async def exercise() -> None:
        fake = FakeWorker(
            (
                _worker_response(request_id="wrong"),
                _worker_response(request_id="gauge-2"),
            )
        )
        operator = AnalogGaugeReaderOperator(
            "reader",
            _params(tmp_path),
            worker_factory=lambda *args, **kwargs: fake,
        )

        failed = (
            await operator.process(_envelope(_request(request_id="gauge-1")))
        ).envelope.payload
        succeeded = (
            await operator.process(_envelope(_request(request_id="gauge-2")))
        ).envelope.payload

        assert failed.status is InferenceStatus.FAILED
        assert failed.error is not None
        assert failed.error.code == "GAUGE_INVALID_RESPONSE"
        assert succeeded.status is InferenceStatus.SUCCEEDED
        health = await operator.health()
        assert health.healthy is True
        assert "completed=1" in health.detail
        assert "failed=1" in health.detail
        await operator.stop()

    asyncio.run(exercise())


def test_worker_failure_response_preserves_error_without_killing_operator(
    tmp_path: Path,
) -> None:
    async def exercise() -> None:
        fake = FakeWorker(
            (
                _worker_response(
                    status="failed",
                    value=None,
                    with_annotated=False,
                    input_image=None,
                ),
                _worker_response(request_id="gauge-2"),
            )
        )
        operator = AnalogGaugeReaderOperator(
            "reader",
            _params(tmp_path),
            worker_factory=lambda *args, **kwargs: fake,
        )

        failed = (
            await operator.process(_envelope(_request(request_id="gauge-1")))
        ).envelope.payload
        succeeded = (
            await operator.process(_envelope(_request(request_id="gauge-2")))
        ).envelope.payload

        assert failed.status is InferenceStatus.FAILED
        assert failed.error is not None
        assert failed.error.code == "GAUGE_READER_FAILED"
        assert failed.error.stage == "ocr"
        assert failed.artifacts == ()
        assert succeeded.status is InferenceStatus.SUCCEEDED
        await operator.stop()

    asyncio.run(exercise())


def test_timeout_becomes_retryable_failed_response(tmp_path: Path) -> None:
    async def exercise() -> None:
        fake = FakeWorker((asyncio.TimeoutError(),))
        operator = AnalogGaugeReaderOperator(
            "reader",
            _params(tmp_path),
            worker_factory=lambda *args, **kwargs: fake,
        )

        response = (
            await operator.process(_envelope(_request(roles=())))
        ).envelope.payload

        assert response.status is InferenceStatus.FAILED
        assert response.error is not None
        assert response.error.code == "GAUGE_TIMEOUT"
        assert response.error.retryable is True
        await operator.stop()

    asyncio.run(exercise())


@pytest.mark.parametrize(
    ("status", "value", "expected"),
    [
        ("no_result", None, InferenceStatus.NO_RESULT),
        ("partial", 3.25, InferenceStatus.PARTIAL),
    ],
)
def test_operator_maps_non_success_statuses(
    tmp_path: Path,
    status: str,
    value: float | None,
    expected: InferenceStatus,
) -> None:
    async def exercise() -> None:
        fake = FakeWorker(
            (
                _worker_response(
                    status=status,
                    value=value,
                    with_annotated=False,
                ),
            )
        )
        operator = AnalogGaugeReaderOperator(
            "reader",
            _params(tmp_path),
            worker_factory=lambda *args, **kwargs: fake,
        )
        response = (
            await operator.process(_envelope(_request(roles=("diagnostics",))))
        ).envelope.payload
        assert response.status is expected
        assert bool(response.outputs) is (value is not None)
        await operator.stop()

    asyncio.run(exercise())


def test_missing_unit_is_partial_with_warning(tmp_path: Path) -> None:
    async def exercise() -> None:
        worker_response = _worker_response(with_annotated=False)
        worker_response = FramedWorkerResponse(
            header={**worker_response.header, "unit": None},
            blobs=worker_response.blobs,
        )
        fake = FakeWorker((worker_response,))
        operator = AnalogGaugeReaderOperator(
            "reader",
            _params(tmp_path),
            worker_factory=lambda *args, **kwargs: fake,
        )
        response = (
            await operator.process(_envelope(_request(roles=())))
        ).envelope.payload
        assert response.status is InferenceStatus.PARTIAL
        assert response.outputs == (
            ScalarOutput(name="reading", value=12.5, unit=None),
        )
        assert "gauge unit was not recognized" in response.warnings
        await operator.stop()

    asyncio.run(exercise())


def test_operator_configuration_is_strict_and_resolves_model_paths(
    tmp_path: Path,
) -> None:
    params = _params(tmp_path)
    fake = FakeWorker((_worker_response(),))
    operator = AnalogGaugeReaderOperator(
        "reader",
        params,
        worker_factory=lambda *args, **kwargs: fake,
    )
    assert operator.params.detection_model_path == (
        tmp_path / "models/detection.pt"
    ).resolve()
    assert operator.params.max_pixels == 1_234_567

    default_params = dict(params)
    default_params.pop("max_pixels")
    default_operator = AnalogGaugeReaderOperator(
        "reader-default-pixels",
        default_params,
        worker_factory=lambda *args, **kwargs: fake,
    )
    assert default_operator.params.max_pixels == 25_000_000

    with pytest.raises(ValidationError, match="Extra inputs are not permitted"):
        AnalogGaugeReaderOperator(
            "reader",
            {**params, "unknown": True},
            worker_factory=lambda *args, **kwargs: fake,
        )
    with pytest.raises(ValidationError, match="absolute path"):
        AnalogGaugeReaderOperator(
            "reader",
            {**params, "python_executable": "python3.8"},
            worker_factory=lambda *args, **kwargs: fake,
        )
    for invalid_max_pixels in (0, 250_000_001):
        with pytest.raises(ValidationError, match="max_pixels"):
            AnalogGaugeReaderOperator(
                "reader",
                {**params, "max_pixels": invalid_max_pixels},
                worker_factory=lambda *args, **kwargs: fake,
            )


def test_operator_command_preserves_virtualenv_python_symlink(
    tmp_path: Path,
) -> None:
    params = _params(tmp_path)
    python_link = tmp_path / "worker-python"
    python_link.symlink_to(sys.executable)
    params["python_executable"] = str(python_link)
    fake = FakeWorker((_worker_response(),))
    constructed: dict[str, object] = {}

    def factory(command, **kwargs):  # type: ignore[no-untyped-def]
        constructed["command"] = tuple(command)
        return fake

    operator = AnalogGaugeReaderOperator(
        "reader",
        params,
        worker_factory=factory,
    )

    command = constructed["command"]
    assert isinstance(command, tuple)
    assert command[0] == str(python_link)
    assert operator.params.python_executable == python_link