app/services/qr_decoder_service.py

from __future__ import annotations

import asyncio
import io
import ipaddress
import os
import socket
import time
from typing import Any, Dict, List, Optional, Tuple
from urllib.parse import urlparse

import cv2
import httpx
import numpy as np
from PIL import Image, UnidentifiedImageError

from app.core.logger import get_logger
from app.core.thread_pool import thread_pool

logger = get_logger(__name__)

try:
    from pyzbar import pyzbar
    PYZBAR_AVAILABLE = True
except Exception:
    PYZBAR_AVAILABLE = False


class QRCodeExtractionError(Exception):
    pass


class QRDecoderResult:
    def __init__(
        self,
        success: bool,
        decoded_data: List[Dict[str, Any]],
        error_message: Optional[str] = None,
        processing_time_ms: float = 0.0,
    ):
        self.success = success
        self.decoded_data = decoded_data
        self.error_message = error_message
        self.processing_time_ms = processing_time_ms


class _PreprocessingPipeline:
    """Collection of static preprocessing strategies for QR code images.
    
    Each strategy returns (preprocessed_image, label) or None if not applicable.
    Strategies are tried in order from fastest/least destructive to most aggressive.
    """

    @staticmethod
    def original(image: np.ndarray) -> Optional[Tuple[np.ndarray, str]]:
        return (image, "original")

    @staticmethod
    def grayscale(image: np.ndarray) -> Optional[Tuple[np.ndarray, str]]:
        if image.shape[2] == 3:
            return (cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), "grayscale")
        return (image, "grayscale")

    @staticmethod
    def otsu_threshold(image: np.ndarray) -> Optional[Tuple[np.ndarray, str]]:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        _, binary = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        return (binary, "otsu")

    @staticmethod
    def adaptive_threshold(image: np.ndarray) -> Optional[Tuple[np.ndarray, str]]:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        binary = cv2.adaptiveThreshold(
            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 51, 2
        )
        return (binary, "adaptive")

    @staticmethod
    def clahe(image: np.ndarray) -> Optional[Tuple[np.ndarray, str]]:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(gray)
        return (enhanced, "clahe")

    @staticmethod
    def unsharp_mask(image: np.ndarray) -> Optional[Tuple[np.ndarray, str]]:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image
        blurred = cv2.GaussianBlur(gray, (0, 0), 3.0)
        sharpened = cv2.addWeighted(gray, 1.5, blurred, -0.5, 0)
        return (sharpened, "unsharp")

    @staticmethod
    def inverted_otsu(image: np.ndarray) -> Optional[Tuple[np.ndarray, str]]:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        _, binary = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        inverted = cv2.bitwise_not(binary)
        return (inverted, "inverted_otsu")

    @staticmethod
    def morphological_clean(image: np.ndarray) -> Optional[Tuple[np.ndarray, str]]:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image
        _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        kernel = np.ones((3, 3), np.uint8)
        cleaned = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_OPEN, kernel)
        return (cleaned, "morphological")

    @classmethod
    def all_strategies(cls) -> List[Any]:
        return [
            cls.original,
            cls.grayscale,
            cls.otsu_threshold,
            cls.adaptive_threshold,
            cls.clahe,
            cls.unsharp_mask,
            cls.inverted_otsu,
            cls.morphological_clean,
        ]


class QRDecoderService:
    def __init__(
        self,
        timeout: float = 30.0,
        max_file_size_mb: float = 20.0,
        allow_private_network_urls: bool = False,
    ) -> None:
        self._timeout = timeout
        self._max_file_size = int(max_file_size_mb * 1024 * 1024)
        self._allow_private_network_urls = allow_private_network_urls
        self._detector = cv2.QRCodeDetector()
        self._preprocessing = _PreprocessingPipeline()

    # ------------------------------------------------------------------
    # Public API  (all public methods are async)
    # ------------------------------------------------------------------

    async def extract(self, source: str) -> QRDecoderResult:
        if self._is_url(source):
            return await self.decode_from_url(source)
        return await self.decode_from_file(source)

    async def decode_from_file(self, file_path: str) -> QRDecoderResult:
        start = time.perf_counter()
        try:
            if not os.path.isfile(file_path):
                elapsed = round((time.perf_counter() - start) * 1000, 3)
                return QRDecoderResult(False, [], f"File not found: {file_path}", elapsed)
            with open(file_path, "rb") as f:
                raw = f.read()
            return await self.decode_from_bytes(raw, source_label=file_path)
        except QRCodeExtractionError as exc:
            elapsed = round((time.perf_counter() - start) * 1000, 3)
            return QRDecoderResult(False, [], str(exc), elapsed)
        except Exception as exc:
            elapsed = round((time.perf_counter() - start) * 1000, 3)
            logger.exception("Unexpected error decoding QR file")
            return QRDecoderResult(False, [], f"Processing failed: {exc}", elapsed)

    async def decode_from_bytes(
        self, image_bytes: bytes, source_label: str = "<bytes>"
    ) -> QRDecoderResult:
        start = time.perf_counter()
        try:
            if len(image_bytes) > self._max_file_size:
                elapsed = round((time.perf_counter() - start) * 1000, 3)
                return QRDecoderResult(
                    False, [],
                    f"Image exceeds maximum size of {self._max_file_size // (1024 * 1024)} MB.",
                    elapsed,
                )

            image = await self._load_cv_image_async(image_bytes, source_label)
            result = await self._decode_async(image, source_label)
            result.processing_time_ms = round((time.perf_counter() - start) * 1000, 3)
            return result
        except QRCodeExtractionError as exc:
            elapsed = round((time.perf_counter() - start) * 1000, 3)
            return QRDecoderResult(False, [], str(exc), elapsed)
        except Exception as exc:
            elapsed = round((time.perf_counter() - start) * 1000, 3)
            logger.exception("Unexpected error decoding QR bytes")
            return QRDecoderResult(False, [], f"Processing failed: {exc}", elapsed)

    async def decode_from_url(self, image_url: str) -> QRDecoderResult:
        start = time.perf_counter()
        try:
            if not self._is_url(image_url):
                elapsed = round((time.perf_counter() - start) * 1000, 3)
                return QRDecoderResult(
                    False, [], f"Not a valid http(s) URL: {image_url}", elapsed
                )

            self._validate_url_is_safe(image_url)
            raw = await self._download(image_url)
            result = await self.decode_from_bytes(raw, source_label=image_url)
            result.processing_time_ms = round((time.perf_counter() - start) * 1000, 3)
            return result
        except QRCodeExtractionError as exc:
            elapsed = round((time.perf_counter() - start) * 1000, 3)
            return QRDecoderResult(False, [], str(exc), elapsed)
        except Exception as exc:
            elapsed = round((time.perf_counter() - start) * 1000, 3)
            logger.exception("Unexpected error decoding QR from URL")
            return QRDecoderResult(False, [], f"Processing failed: {exc}", elapsed)

    # ------------------------------------------------------------------
    # Async wrappers for CPU-bound OpenCV/PIL operations
    # ------------------------------------------------------------------

    async def _load_cv_image_async(self, raw_bytes: bytes, origin: str) -> np.ndarray:
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(thread_pool, self._load_cv_image, raw_bytes, origin)

    async def _decode_async(
        self, image: np.ndarray, source: str
    ) -> QRDecoderResult:
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(thread_pool, self._decode_sync, image, source)

    # ------------------------------------------------------------------
    # Synchronous CPU-bound implementations
    # ------------------------------------------------------------------

    @staticmethod
    def _load_cv_image(raw_bytes: bytes, origin: str) -> np.ndarray:
        try:
            pil_image = Image.open(io.BytesIO(raw_bytes))
            pil_image.load()
        except UnidentifiedImageError as exc:
            raise QRCodeExtractionError(
                f"'{origin}' is not a readable image file"
            ) from exc
        except Exception as exc:
            raise QRCodeExtractionError(
                f"Could not open image '{origin}': {exc}"
            ) from exc

        rgb = pil_image.convert("RGB")
        arr = np.array(rgb)
        return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)

    def _decode_sync(self, image: np.ndarray, source: str) -> QRDecoderResult:
        results = self._decode_with_multi_strategy(image)

        if not results and PYZBAR_AVAILABLE:
            results = self._decode_with_pyzbar(image)

        if not results:
            return QRDecoderResult(
                success=False,
                decoded_data=[],
                error_message=(
                    "No QR code could be detected in the image. "
                    "Make sure the image is clear, in-frame, and not "
                    "excessively skewed or low-resolution."
                ),
            )

        return QRDecoderResult(success=True, decoded_data=results)

    def _decode_with_multi_strategy(
        self, image: np.ndarray
    ) -> List[Dict[str, Any]]:
        tried_strategies: List[str] = []
        seen_data: set = set()

        for strategy in _PreprocessingPipeline.all_strategies():
            processed = strategy(image)
            if processed is None:
                continue
            preprocessed_img, label = processed
            tried_strategies.append(label)

            try:
                ok, decoded_info, points, _ = (
                    self._detector.detectAndDecodeMulti(preprocessed_img)
                )
            except cv2.error:
                ok, decoded_info, points = False, [], None

            if ok:
                results = []
                for i, data in enumerate(decoded_info):
                    if data and data not in seen_data:
                        seen_data.add(data)
                        bbox = (
                            points[i].tolist() if points is not None else None
                        )
                        results.append({
                            "data": data,
                            "type": "QRCODE",
                            "bounding_box": bbox,
                            "decoder": f"opencv_{label}",
                        })
                if results:
                    return results

            try:
                data, points, _ = self._detector.detectAndDecode(preprocessed_img)
            except cv2.error:
                data, points = "", None

            if data and data not in seen_data:
                seen_data.add(data)
                bbox = points.tolist() if points is not None else None
                return [{
                    "data": data,
                    "type": "QRCODE",
                    "bounding_box": bbox,
                    "decoder": f"opencv_{label}",
                }]

        logger.debug("All OpenCV strategies failed: %s", tried_strategies)
        return []

    @staticmethod
    def _decode_with_pyzbar(image: np.ndarray) -> List[Dict[str, Any]]:
        results: List[Dict[str, Any]] = []
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        for obj in pyzbar.decode(gray):
            try:
                data = obj.data.decode("utf-8")
            except UnicodeDecodeError:
                data = obj.data.decode("latin-1", errors="replace")
            bbox = [[p.x, p.y] for p in obj.polygon] if obj.polygon else None
            results.append({
                "data": data,
                "type": obj.type,
                "bounding_box": bbox,
                "decoder": "pyzbar",
            })
        return results

    # ------------------------------------------------------------------
    # URL handling
    # ------------------------------------------------------------------

    @staticmethod
    def _is_url(source: str) -> bool:
        try:
            parsed = urlparse(source)
            return parsed.scheme in ("http", "https") and bool(parsed.netloc)
        except Exception:
            return False

    def _validate_url_is_safe(self, url: str) -> None:
        if self._allow_private_network_urls:
            return

        hostname = urlparse(url).hostname
        if not hostname:
            raise QRCodeExtractionError("URL has no hostname")

        try:
            resolved = socket.getaddrinfo(hostname, None)
        except socket.gaierror as exc:
            raise QRCodeExtractionError(
                f"Could not resolve host '{hostname}': {exc}"
            ) from exc

        for family, _, _, _, sockaddr in resolved:
            ip_str = sockaddr[0]
            try:
                ip_obj = ipaddress.ip_address(ip_str)
            except ValueError:
                continue
            if (
                ip_obj.is_private
                or ip_obj.is_loopback
                or ip_obj.is_link_local
                or ip_obj.is_reserved
            ):
                raise QRCodeExtractionError(
                    f"Refusing to fetch URL: host resolves to a "
                    f"non-public address ({ip_str})"
                )

    async def _download(self, url: str) -> bytes:
        try:
            async with httpx.AsyncClient(
                timeout=self._timeout, follow_redirects=True
            ) as client:
                async with client.stream("GET", url) as resp:
                    resp.raise_for_status()

                    content_length = resp.headers.get("Content-Length")
                    if content_length is not None:
                        try:
                            if int(content_length) > self._max_file_size:
                                raise QRCodeExtractionError(
                                    f"Remote file too large "
                                    f"(Content-Length={content_length} bytes)"
                                )
                        except ValueError:
                            pass

                    chunks = []
                    total = 0
                    async for chunk in resp.aiter_bytes(chunk_size=65536):
                        total += len(chunk)
                        if total > self._max_file_size:
                            raise QRCodeExtractionError(
                                f"Download exceeded max allowed size "
                                f"of {self._max_file_size} bytes"
                            )
                        chunks.append(chunk)

                    data = b"".join(chunks)
                    if not data:
                        raise QRCodeExtractionError(
                            "Downloaded content was empty"
                        )
                    return data
        except httpx.HTTPStatusError as exc:
            raise QRCodeExtractionError(
                f"Failed to fetch image from URL: HTTP {exc.response.status_code}"
            ) from exc
        except httpx.RequestError as exc:
            raise QRCodeExtractionError(
                f"Failed to fetch image from URL: {exc}"
            ) from exc