import io

import cv2
import numpy as np
from aarelcinfer_client import AuthenticatedClient
from aarelcinfer_client.api import beam, config, predictions
from aarelcinfer_client.models import LatestPredictionModel, RuntimeConfigPatchModel
from PIL import Image

from aarecommon.config.beamline import cfg_get, mx_beamline
from aarecommon.models.beamline import MXBeamline


class AareLCInferWrapper:
    def __init__(self, bl: MXBeamline, secret: str = "1s3ng@rd"):
        if bl == MXBeamline.X10SA or bl == MXBeamline.X06DA:
            host = cfg_get("daq.hardware.aarelc_url")
            if host is None:
                raise ValueError("AareLCInferWrapper: AareLC URL not configured")
        elif bl == MXBeamline.X06SA or bl == MXBeamline.SIMULATED:
            raise NotImplementedError(f"AareLCInferWrapper not implemented for {bl}")
        else:
            raise ValueError(f"Unknown beamline {bl}")

        self.client = AuthenticatedClient(base_url=host, api_key=secret)
        self.client.headers["X-API-Key"] = secret
        self._host = host
        self._api_config = config
        self._api_predictions = predictions
        self._api_beam = beam

    def get_config(self) -> config.ConfigSnapshotResponse:
        return self._api_config.get_config(self.client)

    def update_config(self, patch: RuntimeConfigPatchModel) -> config.ConfigUpdateResponse:
        return self._api_config.update_config(self.client, patch)

    def get_latest_prediction(self) -> LatestPredictionModel:
        return self._api_predictions.get_latest_prediction(self.client)

    def get_latest_frame_png(self):
        return self._api_predictions.get_latest_frame_png(self.client)

    def get_latest_prediction_bundle(self):
        return self._api_predictions.get_latest_prediction_bundle(self.client)

    def send_samcam_details(self, beam_mark, beam_dimensions):
        return self._api_beam.set_beam_mark(beam_mark, beam_dimensions)


if __name__ == "__main__":
    wrapper = AareLCInferWrapper(bl=mx_beamline())

    try:
        print("Fetching config...")
        config = wrapper.get_config()
        print(f"Config: {config}")
        print("Updating config...")
        patch = RuntimeConfigPatchModel(
            conf=0.4
            # infer_scale: float | None = Field(default=None, gt=0.0, le=1.0)
            # skip: int | None = Field(default=None, ge=0)
            # max_fps: float | None = Field(default=None, ge=0.0)
            # device: str | None = None
            # publish_enabled: bool | None = None
            # compute_target_point: bool | None = None
            # focus_enabled: bool | None = None
            # focus_epics_enabled: bool | None = None
            # focus_pv: str | None = None
            # focus_every: int | None = Field(default=None, ge=1)
            # focus_scale: float | None = Field(default=None, gt=0.0, le=1.0)
            # focus_pv_min_period_ms: float | None = Field(default=None, ge=0.0)
            # tracker: Literal["none", "bytetrack", "botsort"] | None = None
            # pt: str | None = None
            # engine: str | None = None
            # task: Literal["auto", "detect", "segment"] | None = None
        )
        response = wrapper.update_config(patch)
        print("Config updated!")
        print(f"Model reloaded: {getattr(response, 'model_reloaded', False)}")
    except Exception as e:  # noqa
        print(f"Error updating config: {e}")

    try:
        print("Fetching bundle...")
        bundle = wrapper.get_latest_prediction_bundle()

        jpeg_bytes = bundle.image_jpeg

        # The client seems to already return the validated model as 'metadata'
        prediction = bundle.metadata

        # Safety check: if it's still a dict for some reason, validate it; otherwise use as is
        if isinstance(prediction, dict):
            prediction = LatestPredictionModel.model_validate(prediction)

        print(
            f"Received bundle. Image: {len(jpeg_bytes)} bytes. Detections: {len(prediction.boxes)}"
        )

        # Decode Image
        img = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB")
        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

        # Draw Overlay using the model attributes
        print(prediction.target_point.x, prediction.target_point.y)
        print("boxes: ", prediction.boxes)
        for det in prediction.boxes:
            x1, y1 = int(det.x1), int(det.y1)
            x2, y2 = int(det.x2), int(det.y2)

            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(
                img,
                f"{det.label} {det.conf:.2f}",
                (x1, max(20, y1 - 8)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6,
                (0, 255, 0),
                2,
            )

            if det.poly:
                # Convert to numpy and shift coordinates from relative to absolute
                pts = np.array(det.poly, dtype=np.int32)
                pts[:, 0] += x1  # Shift X
                pts[:, 1] += y1  # Shift Y

                pts = pts.reshape((-1, 1, 2))
                cv2.polylines(img, [pts], isClosed=True, color=(0, 0, 255), thickness=2)

        # cv2.imshow("Latest Prediction Bundle", img)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()

    except Exception as e:  # noqa
        print(f"Error processing prediction bundle: {e}")