"""Document validators — MIME detection and file safety checks.

Validates uploaded documents before any processing:
1. File size limit
2. MIME type detection via magic bytes (content-first, not extension-first)
3. Supported MIME type check
4. Basic structural sanity

Never calls AI models.
Never crashes on corrupt input — returns structured errors.
"""

from __future__ import annotations

import hashlib
import json
import logging

from ehr_intel.models import DocumentTypeEnum

log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Supported MIME types
# ---------------------------------------------------------------------------

SUPPORTED_MIME_TYPES: frozenset[str] = frozenset({
    "application/pdf",
    "image/png",
    "image/jpeg",
    "image/webp",
    "text/plain",
    "application/json",
    "application/fhir+json",
})

# Magic byte signatures for content-based MIME detection
_MAGIC_SIGNATURES: list[tuple[bytes, str]] = [
    (b"%PDF", "application/pdf"),
    (b"\x89PNG\r\n\x1a\n", "image/png"),
    (b"\xff\xd8\xff", "image/jpeg"),
    (b"RIFF", "image/webp"),   # bytes 8-11 must also be b"WEBP"
    (b"\xff\xd8", "image/jpeg"),  # looser JPEG prefix fallback
]

_EXTENSION_MIME_MAP: dict[str, str] = {
    ".pdf": "application/pdf",
    ".png": "image/png",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".webp": "image/webp",
    ".txt": "text/plain",
    ".json": "application/json",
}


class ValidationError(Exception):
    """Raised when a document fails validation."""

    def __init__(self, message: str, code: str = "VALIDATION_ERROR") -> None:
        super().__init__(message)
        self.code = code
        self.message = message


class ValidationResult:
    """Result of document validation."""

    def __init__(
        self,
        *,
        valid: bool,
        mime_type: str,
        document_type: DocumentTypeEnum,
        file_size_bytes: int,
        sha256: str,
        error: str | None = None,
        error_code: str | None = None,
    ) -> None:
        self.valid = valid
        self.mime_type = mime_type
        self.document_type = document_type
        self.file_size_bytes = file_size_bytes
        self.sha256 = sha256
        self.error = error
        self.error_code = error_code


def detect_mime_type(
    raw_bytes: bytes,
    filename: str | None = None,
    claimed_mime: str | None = None,
) -> str:
    """Detects MIME type using magic bytes first.

    Falls back to claimed MIME and then file extension.
    Content always wins over extension claim.

    Args:
        raw_bytes: Raw file bytes. May be empty.
        filename: Original filename. Used only as final fallback.
        claimed_mime: MIME type claimed by the caller. Not trusted blindly.

    Returns:
        Detected MIME type string.
    """
    if raw_bytes:
        # --- Primary: magic byte sniffing ---
        for signature, mime in _MAGIC_SIGNATURES:
            if raw_bytes.startswith(signature):
                if signature == b"RIFF":
                    # WebP: bytes 8-12 must be b"WEBP"
                    if len(raw_bytes) >= 12 and raw_bytes[8:12] == b"WEBP":
                        return "image/webp"
                    continue
                return mime

        # --- Check for JSON / FHIR Bundle ---
        try:
            snippet = raw_bytes[:1024].decode("utf-8", errors="ignore").strip()
            if snippet.startswith("{") or snippet.startswith("["):
                parsed = json.loads(raw_bytes.decode("utf-8", errors="replace"))
                if isinstance(parsed, dict) and parsed.get("resourceType") == "Bundle":
                    return "application/fhir+json"
                return "application/json"
        except (json.JSONDecodeError, UnicodeDecodeError):
            pass

        # --- Check for plain text (UTF-8 decodable, no binary nulls) ---
        try:
            text = raw_bytes[:512].decode("utf-8")
            if "\x00" not in text:
                # Looks like plain text
                if claimed_mime == "text/plain" or (
                    filename and filename.lower().endswith(".txt")
                ):
                    return "text/plain"
        except UnicodeDecodeError:
            pass

    # --- Secondary: claimed MIME ---
    if claimed_mime:
        normalized = claimed_mime.lower().strip()
        # Normalize "image/jpg" → "image/jpeg"
        if normalized == "image/jpg":
            normalized = "image/jpeg"
        if normalized in SUPPORTED_MIME_TYPES or normalized == "text/plain":
            return normalized

    # --- Tertiary: extension ---
    if filename:
        fn_lower = filename.lower().strip()
        for ext, mime in _EXTENSION_MIME_MAP.items():
            if fn_lower.endswith(ext):
                return mime

    return "application/octet-stream"


def classify_document_type(
    *,
    mime_type: str,
    filename: str | None = None,
    text_content: str | None = None,
) -> DocumentTypeEnum:
    """Classifies clinical document type from available signals.

    This is a lightweight heuristic classification. Gemini will refine it later.

    Args:
        mime_type: Detected MIME type.
        filename: Original filename.
        text_content: Short text sample, if available (e.g. from PDF extraction).

    Returns:
        DocumentTypeEnum best guess.
    """
    if mime_type == "application/fhir+json":
        return DocumentTypeEnum.FHIR_BUNDLE

    probe = " ".join(filter(None, [filename or "", text_content or ""])).lower()

    # FHIR Bundle text indicator
    if "resourcetype" in probe or '"bundle"' in probe:
        return DocumentTypeEnum.FHIR_BUNDLE

    # Prescription indicators
    rx_kws = ["prescription", "rx:", "sig:", "disp:", "dispense", "refill"]
    rx_soft = ["tab", "cap", "od", "bd", "tid", "qid", "po", "mg", "tablet", "capsule"]
    if any(kw in probe for kw in rx_kws) or sum(1 for kw in rx_soft if kw in probe) >= 3:
        if "discharge summary" not in probe:
            return DocumentTypeEnum.PRESCRIPTION

    # Doctor note indicators
    note_kws = ["progress note", "soap note", "subjective:", "objective:", "assessment", "consultation note"]
    if any(kw in probe for kw in note_kws):
        return DocumentTypeEnum.DOCTOR_NOTE

    # Discharge summary indicators
    ds_kws = ["discharge summary", "hospital course", "admission date", "discharge date", "discharge diagnosis"]
    if any(kw in probe for kw in ds_kws):
        return DocumentTypeEnum.DISCHARGE_SUMMARY

    # Lab report indicators
    lab_kws = ["lab report", "laboratory", "reference range", "specimen", "hemoglobin", "creatinine"]
    if any(kw in probe for kw in lab_kws):
        return DocumentTypeEnum.LAB_REPORT

    # Filename-based hints
    if filename:
        fn = filename.lower()
        if any(kw in fn for kw in ["presc", "rx_", "_rx"]):
            return DocumentTypeEnum.PRESCRIPTION
        if any(kw in fn for kw in ["discharge", "disch", "summary"]):
            return DocumentTypeEnum.DISCHARGE_SUMMARY
        if any(kw in fn for kw in ["note", "soap", "doctor"]):
            return DocumentTypeEnum.DOCTOR_NOTE
        if "lab" in fn or "report" in fn:
            return DocumentTypeEnum.LAB_REPORT

    if mime_type in ("image/png", "image/jpeg", "image/webp"):
        return DocumentTypeEnum.GENERAL_CLINICAL

    return DocumentTypeEnum.UNKNOWN


def validate_document(
    raw_bytes: bytes,
    *,
    filename: str | None = None,
    claimed_mime: str | None = None,
    max_file_mb: int = 20,
) -> ValidationResult:
    """Full document validation pipeline.

    Args:
        raw_bytes: Document bytes (may be empty — will fail gracefully).
        filename: Original filename.
        claimed_mime: MIME type claimed by caller.
        max_file_mb: Maximum allowed file size in MB.

    Returns:
        ValidationResult — always returns (never raises), callers check .valid.
    """
    # Compute hash immediately (needed for logging even on failure)
    sha256 = hashlib.sha256(raw_bytes).hexdigest() if raw_bytes else "empty"
    size = len(raw_bytes)

    # --- Empty file check ---
    if not raw_bytes or size == 0:
        return ValidationResult(
            valid=False,
            mime_type="application/octet-stream",
            document_type=DocumentTypeEnum.UNKNOWN,
            file_size_bytes=0,
            sha256=sha256,
            error="Empty file received.",
            error_code="EMPTY_FILE",
        )

    # --- Size limit check ---
    max_bytes = max_file_mb * 1024 * 1024
    if size > max_bytes:
        return ValidationResult(
            valid=False,
            mime_type="application/octet-stream",
            document_type=DocumentTypeEnum.UNKNOWN,
            file_size_bytes=size,
            sha256=sha256,
            error=f"File size {size / 1024 / 1024:.1f} MB exceeds limit of {max_file_mb} MB.",
            error_code="FILE_TOO_LARGE",
        )

    # --- MIME detection ---
    mime = detect_mime_type(raw_bytes, filename=filename, claimed_mime=claimed_mime)

    # --- Supported MIME check ---
    if mime not in SUPPORTED_MIME_TYPES and mime != "application/octet-stream":
        # Allow some unlisted mimes through with a warning
        pass

    if mime == "application/octet-stream":
        return ValidationResult(
            valid=False,
            mime_type=mime,
            document_type=DocumentTypeEnum.UNKNOWN,
            file_size_bytes=size,
            sha256=sha256,
            error=(
                f"Unsupported or unrecognized file type. "
                f"Accepted: {sorted(SUPPORTED_MIME_TYPES)}."
            ),
            error_code="UNSUPPORTED_MIME",
        )

    # --- Document type classification ---
    doc_type = classify_document_type(mime_type=mime, filename=filename)

    log.info(
        "Document validated: sha256=%s... size=%d mime=%s type=%s",
        sha256[:8],
        size,
        mime,
        doc_type,
    )

    return ValidationResult(
        valid=True,
        mime_type=mime,
        document_type=doc_type,
        file_size_bytes=size,
        sha256=sha256,
    )


__all__ = [
    "SUPPORTED_MIME_TYPES",
    "ValidationError",
    "ValidationResult",
    "classify_document_type",
    "detect_mime_type",
    "validate_document",
]