import io
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from upf.config import UPFConfig
from upf.detectors.image_layout import LayoutSignal, analyse_layout
from upf.pipeline.text_pipeline import TextPipeline
from upf.redaction.image import apply_face_blur, apply_image_redaction, apply_signature_blur, build_word_boxes

# Risk boost applied to risk_score when a layout signal is detected.
_LAYOUT_RISK_BOOST: Dict[str, int] = {
    "mrz": 10,    # passport / travel document — very high sensitivity
    "id_card": 7, # national ID card
    "table": 4,   # bank statement / invoice
}

# Layout types that force risk_level = 'high' regardless of text risk
_LAYOUT_HIGH_RISK_TYPES = frozenset({"mrz", "id_card"})


@dataclass
class ImageSanitizeResult:
    sanitized_text: str
    sanitized_image_bytes: bytes
    entities: List[Dict[str, Any]]
    risk_score: int = 0
    risk_level: str = "low"
    layout_signals: List[LayoutSignal] = field(default_factory=list)


class ImagePipeline:
    MAX_IMAGE_SIZE_BYTES = 10 * 1024 * 1024

    def __init__(self, config: Optional[UPFConfig] = None) -> None:
        self._text_pipeline = TextPipeline(config)
        self._config = config or UPFConfig()

    def sanitize_ocr_text(self, ocr_text: str) -> str:
        return self._text_pipeline.sanitize(ocr_text)

    def sanitize_image_bytes(
        self, image_bytes: bytes, ocr_text: Optional[str] = None
    ) -> ImageSanitizeResult:
        image = self.load_image(image_bytes)
        extracted_text = ocr_text if ocr_text is not None else self.extract_text(image)
        # Align word-box char offsets to the exact OCR text used for entity detection.
        word_boxes = build_word_boxes(image, ocr_text=extracted_text)

        # Text-based detection
        text_result = self._text_pipeline.sanitize_with_details(extracted_text)

        # Layout analysis
        w, h = image.size
        layout_signals = analyse_layout(word_boxes, extracted_text, w, h)

        # Risk aggregation: text score + layout boost
        layout_boost = _compute_layout_boost(layout_signals)
        total_risk_score = text_result.risk_score + layout_boost
        risk_level = _aggregate_risk_level(text_result.risk_level, layout_signals)

        # Visual redaction: text entities only (values, not labels).
        # Layout pseudo-entities are intentionally excluded — the text detectors
        # already redact the PII values inside MRZ/ID-card regions, and a single
        # min-start/max-end span would black out the entire document section.
        redacted_image = apply_image_redaction(image, text_result.entities, word_boxes=word_boxes)
        # Face blur: enabled by default (config.blur_faces=True) for ID card portraits.
        redacted_image = apply_face_blur(
            redacted_image,
            enabled=self._config.blur_faces,
            blur_strength=self._config.face_blur_strength,
        )
        # Signature blur: detect on the original clean image so that black
        # redaction patches don't corrupt Otsu thresholding, then composite
        # the blurred regions onto the already-redacted image.
        redacted_image = apply_signature_blur(
            redacted_image,
            enabled=self._config.blur_signatures,
            blur_strength=self._config.signature_blur_strength,
            word_boxes=word_boxes,
            source_image=image,
        )

        return ImageSanitizeResult(
            sanitized_text=text_result.sanitized_text,
            sanitized_image_bytes=self.encode_image_png(redacted_image),
            entities=text_result.entities,
            risk_score=total_risk_score,
            risk_level=risk_level,
            layout_signals=layout_signals,
        )

    def load_image(self, image_bytes: bytes):
        if len(image_bytes) > self.MAX_IMAGE_SIZE_BYTES:
            raise ValueError("Image too large")
        if not image_bytes:
            raise ValueError("Empty image payload")

        try:
            from PIL import Image
        except ImportError as exc:  # pragma: no cover
            raise RuntimeError("Pillow is required for image byte processing") from exc

        image = Image.open(io.BytesIO(image_bytes))
        return image.convert("RGB")

    def extract_text(self, image) -> str:
        try:
            import pytesseract
        except ImportError as exc:
            raise RuntimeError(
                "OCR backend unavailable. Install pytesseract or provide ocr_text."
            ) from exc
        return pytesseract.image_to_string(image).strip()

    def encode_image_png(self, image) -> bytes:
        buf = io.BytesIO()
        image.save(buf, format="PNG")
        return buf.getvalue()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _compute_layout_boost(signals: List[LayoutSignal]) -> int:
    boost = 0
    for sig in signals:
        base = _LAYOUT_RISK_BOOST.get(sig.layout_type, 0)
        boost += int(base * sig.confidence)
    return boost


def _aggregate_risk_level(text_level: str, signals: List[LayoutSignal]) -> str:
    if text_level == "high":
        return "high"
    for sig in signals:
        if sig.layout_type in _LAYOUT_HIGH_RISK_TYPES and sig.confidence >= 0.6:
            return "high"
        if sig.layout_type == "table" and sig.confidence >= 0.7 and text_level == "low":
            return "medium"
    return text_level


def _layout_signals_to_entities(signals: List[LayoutSignal]) -> List[Dict[str, Any]]:
    """Convert high-confidence layout signals to pseudo-entities for visual redaction."""
    entities = []
    for sig in signals:
        if sig.layout_type not in _LAYOUT_HIGH_RISK_TYPES:
            continue
        if sig.confidence < 0.70:
            continue
        if not sig.ocr_text_spans:
            continue
        min_start = min(s for s, _ in sig.ocr_text_spans)
        max_end = max(e for _, e in sig.ocr_text_spans)
        entities.append({
            "type": sig.layout_type.upper(),
            "start": min_start,
            "end": max_end,
            "text": "",
        })
    return entities
