database.py

import hashlib
import json
import logging
import os
import re
from datetime import datetime, timezone
from typing import Any, Dict, Optional

import requests


logger = logging.getLogger("clinical-ai-db")

SUPABASE_URL = os.environ.get("SUPABASE_URL", "").rstrip("/")
SUPABASE_SERVICE_ROLE_KEY = os.environ.get("SUPABASE_SERVICE_ROLE_KEY", "")

# The production Supabase bucket used by the working Kaggle server and the
# existing public URLs is exactly `xray-images`.  Do not silently switch to a
# different bucket name.  A stale HF secret must be diagnosed, not hidden.
SUPABASE_BUCKET = (
    os.environ.get("SUPABASE_STORAGE_BUCKET", "xray-images").strip()
    or "xray-images"
)

PUBLIC_STORAGE = (
    os.environ.get("SUPABASE_STORAGE_PUBLIC", "true").strip().lower() == "true"
)

REQUEST_TIMEOUT = 60
UPLOAD_TIMEOUT = 120


def enabled() -> bool:
    return bool(SUPABASE_URL and SUPABASE_SERVICE_ROLE_KEY)


def config_info() -> Dict[str, Any]:
    return {
        "configured": enabled(),
        "url_configured": bool(SUPABASE_URL),
        "service_role_key_configured": bool(SUPABASE_SERVICE_ROLE_KEY),
        "bucket": SUPABASE_BUCKET,
        "public_storage": PUBLIC_STORAGE,
    }


def _safe_response_text(response: requests.Response, limit: int = 1500) -> str:
    try:
        text = response.text or ""
    except Exception:
        text = "<unable to read response body>"
    return text[:limit]


def _list_buckets() -> list:
    """Return the buckets visible to the server-side Supabase key."""
    url = f"{SUPABASE_URL}/storage/v1/bucket"
    response = requests.get(
        url,
        headers=_headers(),
        timeout=REQUEST_TIMEOUT,
    )
    if response.status_code >= 300:
        raise RuntimeError(
            f"Storage bucket list failed ({response.status_code}): "
            f"{_safe_response_text(response, 1200)}"
        )
    try:
        data = response.json() if response.text else []
    except ValueError as exc:
        raise RuntimeError(
            "Supabase returned invalid JSON while listing Storage buckets."
        ) from exc
    return data if isinstance(data, list) else []


def _ensure_bucket() -> None:
    """
    Verify/create the exact configured Storage bucket.

    This deliberately follows the known-working Kaggle behavior: first check
    the exact bucket, then create it if missing, then verify it again.  It does
    not silently rename or switch buckets.  This makes Hugging Face failures
    actionable instead of masking a bad secret/project configuration.
    """
    if not enabled():
        raise RuntimeError("Supabase is not configured on the server.")

    bucket = SUPABASE_BUCKET
    base_url = f"{SUPABASE_URL}/storage/v1/bucket"
    bucket_url = f"{base_url}/{bucket}"

    logger.info("STORAGE: bucket preflight name=%s", bucket)

    # --------------------------------------------------------
    # 1. Direct check of the exact bucket
    # --------------------------------------------------------
    try:
        check = requests.get(
            bucket_url,
            headers=_headers(),
            timeout=REQUEST_TIMEOUT,
        )
    except requests.RequestException as exc:
        raise RuntimeError(
            f"Storage bucket check failed for '{bucket}': {exc}"
        ) from exc

    logger.info(
        "STORAGE: bucket GET name=%s status=%s body=%s",
        bucket,
        check.status_code,
        _safe_response_text(check, 500),
    )

    if check.status_code == 200:
        logger.info("STORAGE: bucket verified name=%s", bucket)
        return

    # --------------------------------------------------------
    # 2. If 404, inspect the actual bucket list for diagnostics.
    # --------------------------------------------------------
    listed_names = []
    if check.status_code == 404:
        try:
            listed = _list_buckets()
            listed_names = [
                str(item.get("name") or item.get("id") or "").strip()
                for item in listed
                if isinstance(item, dict)
            ]
            listed_names = [x for x in listed_names if x]
            logger.warning(
                "STORAGE: requested bucket '%s' not found. Visible buckets=%s",
                bucket,
                listed_names,
            )
        except Exception as exc:
            logger.warning(
                "STORAGE: bucket list diagnostic failed after 404: %s",
                exc,
            )

    # Any non-404 error is a real authentication/permission/API problem; do
    # not try to create another bucket and hide it.
    if check.status_code not in (404,):
        raise RuntimeError(
            f"Storage bucket '{bucket}' check failed: "
            f"{check.status_code} {_safe_response_text(check, 1500)}"
        )

    # --------------------------------------------------------
    # 3. Bucket is absent: create the exact configured bucket.
    # --------------------------------------------------------
    logger.warning(
        "STORAGE: bucket '%s' is missing; attempting creation",
        bucket,
    )

    payload = {
        "id": bucket,
        "name": bucket,
        "public": PUBLIC_STORAGE,
        "file_size_limit": 15728640,
        "allowed_mime_types": [
            "image/png",
            "image/jpeg",
            "image/webp",
            "image/bmp",
        ],
    }

    try:
        created = requests.post(
            base_url,
            headers=_headers(),
            json=payload,
            timeout=REQUEST_TIMEOUT,
        )
    except requests.RequestException as exc:
        raise RuntimeError(
            f"Storage bucket creation failed for '{bucket}': {exc}"
        ) from exc

    logger.info(
        "STORAGE: bucket CREATE name=%s status=%s body=%s",
        bucket,
        created.status_code,
        _safe_response_text(created, 1000),
    )

    if created.status_code not in (200, 201, 409):
        raise RuntimeError(
            f"Storage bucket '{bucket}' could not be created: "
            f"{created.status_code} {_safe_response_text(created, 1500)}"
        )

    # --------------------------------------------------------
    # 4. Final verification
    # --------------------------------------------------------
    try:
        verify = requests.get(
            bucket_url,
            headers=_headers(),
            timeout=REQUEST_TIMEOUT,
        )
    except requests.RequestException as exc:
        raise RuntimeError(
            f"Storage bucket final verification failed for '{bucket}': {exc}"
        ) from exc

    logger.info(
        "STORAGE: bucket VERIFY name=%s status=%s body=%s",
        bucket,
        verify.status_code,
        _safe_response_text(verify, 800),
    )

    if verify.status_code != 200:
        raise RuntimeError(
            "Storage bucket was not available after creation attempt. "
            f"bucket='{bucket}', create_status={created.status_code}, "
            f"verify_status={verify.status_code}, "
            f"verify_response={_safe_response_text(verify, 1200)}, "
            f"visible_buckets={listed_names}"
        )

    logger.info("STORAGE: bucket ready name=%s", bucket)

def _headers(content_type: str = "application/json") -> Dict[str, str]:
    if not enabled():
        raise RuntimeError("Supabase is not configured on the server.")

    return {
        "apikey": SUPABASE_SERVICE_ROLE_KEY,
        "Authorization": f"Bearer {SUPABASE_SERVICE_ROLE_KEY}",
        "Content-Type": content_type,
    }


def _rest_url(table: str) -> str:
    return f"{SUPABASE_URL}/rest/v1/{table}"


def _storage_url(path: str) -> str:
    return f"{SUPABASE_URL}/storage/v1/object/{SUPABASE_BUCKET}/{path}"


def _safe_response_text(response: requests.Response, limit: int = 1500) -> str:
    try:
        text = response.text or ""
    except Exception:
        text = "<unable to read response body>"
    return text[:limit]


def upload_bytes(data: bytes, path: str, content_type: str) -> str:
    """Upload one object to Supabase Storage and return its URL."""
    if not enabled():
        raise RuntimeError("Supabase is not configured on the server.")

    if not data:
        raise ValueError(f"Cannot upload empty data: {path}")

    _ensure_bucket()
    url = _storage_url(path)

    logger.info(
        "STORAGE: uploading path=%s bytes=%d content_type=%s",
        path,
        len(data),
        content_type,
    )

    try:
        response = requests.post(
            url,
            headers={
                **_headers(content_type),
                "x-upsert": "true",
                "cache-control": "3600",
            },
            data=data,
            timeout=UPLOAD_TIMEOUT,
        )
    except requests.RequestException as exc:
        logger.exception("STORAGE: network error path=%s", path)
        raise RuntimeError(
            f"Storage request failed for '{path}': {exc}"
        ) from exc

    if response.status_code not in (200, 201):
        body = _safe_response_text(response)
        logger.error(
            "STORAGE: failed path=%s status=%s body=%s",
            path,
            response.status_code,
            body,
        )
        raise RuntimeError(
            f"Storage upload failed ({response.status_code}) for '{path}': {body}"
        )

    logger.info(
        "STORAGE: upload OK path=%s status=%s",
        path,
        response.status_code,
    )

    if PUBLIC_STORAGE:
        return (
            f"{SUPABASE_URL}/storage/v1/object/public/"
            f"{SUPABASE_BUCKET}/{path}"
        )

    return f"{SUPABASE_URL}/storage/v1/object/{SUPABASE_BUCKET}/{path}"


def _json_request(
    method: str,
    table: str,
    payload: Optional[Dict[str, Any]] = None,
    params: Optional[Dict[str, str]] = None,
):
    """Call Supabase REST with detailed, safe diagnostics."""
    if not enabled():
        raise RuntimeError("Supabase is not configured on the server.")

    url = _rest_url(table)

    logger.info(
        "SUPABASE: %s table=%s params=%s",
        method,
        table,
        params,
    )

    try:
        response = requests.request(
            method,
            url,
            headers={
                **_headers(),
                "Prefer": "return=representation,resolution=merge-duplicates",
            },
            json=payload,
            params=params,
            timeout=REQUEST_TIMEOUT,
        )
    except requests.RequestException as exc:
        logger.exception("SUPABASE: network error method=%s table=%s", method, table)
        raise RuntimeError(
            f"Supabase REST request failed ({method} {table}): {exc}"
        ) from exc

    if response.status_code >= 300:
        body = _safe_response_text(response)
        logger.error(
            "SUPABASE: request failed method=%s table=%s status=%s body=%s",
            method,
            table,
            response.status_code,
            body,
        )
        raise RuntimeError(
            f"Supabase REST {method} failed ({response.status_code}) "
            f"for table '{table}': {body}"
        )

    if not response.text:
        return []

    try:
        result = response.json()
    except ValueError as exc:
        logger.exception("SUPABASE: invalid JSON response table=%s", table)
        raise RuntimeError(
            f"Supabase returned invalid JSON for table '{table}'."
        ) from exc

    logger.info(
        "SUPABASE: request OK method=%s table=%s status=%s",
        method,
        table,
        response.status_code,
    )
    return result


def _localized_patient(patient: Dict[str, Any]) -> Dict[str, Any]:
    gender = str(patient.get("gender") or "")

    gender_ar = {
        "male": "ذكر",
        "female": "أنثى",
    }.get(gender.lower(), gender)

    gender_en = {
        "ذكر": "Male",
        "أنثى": "Female",
    }.get(gender, gender)

    return {
        "name": patient.get("patient_name", ""),
        "birth_date": patient.get("birth_date", ""),
        "age": patient.get("age"),
        "gender": {
            "value": gender,
            "ar": gender_ar,
            "en": gender_en,
        },
    }


def _visualization_bytes(value: Any, key: str) -> bytes:
    """Accept only binary visualization data. Base64 is deliberately unsupported."""
    if isinstance(value, bytes):
        if not value:
            raise ValueError(f"Empty visualization data for '{key}'.")
        return value
    if isinstance(value, bytearray):
        data = bytes(value)
        if not data:
            raise ValueError(f"Empty visualization data for '{key}'.")
        return data
    raise TypeError(
        f"Visualization '{key}' must be binary bytes; Base64/string payloads are not supported."
    )


def save_analysis(
    *,
    result: Dict[str, Any],
    original_bytes: bytes,
    original_filename: str,
    patient: Dict[str, Any],
    user: Dict[str, Any],
) -> Dict[str, Any]:
    """
    Persist a completed inference result.

    This function raises on Storage/REST failure so app.py can report the exact
    database failure. app.py deliberately treats persistence failure as
    non-fatal for the AI result and returns database.saved=false.
    """
    if not enabled():
        logger.warning("DB: Supabase is not configured; persistence skipped.")
        return {
            "saved": False,
            "reason": "database_not_configured",
        }

    analysis_id = str(result.get("analysis_id") or "").strip()
    if not analysis_id:
        raise ValueError("Missing analysis_id in inference result.")

    if not original_bytes:
        raise ValueError("Original image bytes are empty.")

    patient = patient or {}
    user = user or {}

    owner_id = str(user.get("user_id") or "anonymous").strip() or "anonymous"
    folder = f"{owner_id}/{analysis_id}"

    extension = (
        original_filename.rsplit(".", 1)[-1].lower()
        if "." in original_filename
        else "jpg"
    )

    if extension not in {"jpg", "jpeg", "png", "webp", "bmp"}:
        extension = "jpg"

    content_type = {
        "jpg": "image/jpeg",
        "jpeg": "image/jpeg",
        "png": "image/png",
        "webp": "image/webp",
        "bmp": "image/bmp",
    }[extension]

    logger.info(
        "DB: BEGIN save_analysis analysis_id=%s folder=%s user_id=%s",
        analysis_id,
        folder,
        user.get("user_id"),
    )

    # ------------------------------------------------------------
    # 1) Original image
    # ------------------------------------------------------------
    original_path = f"{folder}/original.{extension}"
    original_url = upload_bytes(
        original_bytes,
        original_path,
        content_type,
    )

    # ------------------------------------------------------------
    # 2) Visualizations
    # ------------------------------------------------------------
    visual_urls: Dict[str, str] = {}

    visualization_files = {
        "heatmap_enhanced": ("heatmap.png", "image/png"),
        "overlay": ("overlay.jpg", "image/jpeg"),
        "contour": ("contour.jpg", "image/jpeg"),
        "threshold": ("threshold.jpg", "image/jpeg"),
        "comparison": ("comparison.jpg", "image/jpeg"),
    }

    visualizations = result.get("visualizations") or {}

    for key, (name, mime_type) in visualization_files.items():
        value = visualizations.get(key) or ""

        if not value:
            logger.warning(
                "DB: visualization missing; skipping key=%s",
                key,
            )
            continue

        binary = _visualization_bytes(value, key)
        visual_path = f"{folder}/{name}"

        visual_urls[key] = upload_bytes(
            binary,
            visual_path,
            mime_type,
        )

    # ------------------------------------------------------------
    # 3) Prepare database row
    # ------------------------------------------------------------
    reports_ar = result.get("medical_report_ar")
    reports_en = result.get("medical_report_en")

    if not isinstance(reports_ar, dict) or not isinstance(reports_en, dict):
        raise ValueError(
            "Inference result does not contain both medical_report_ar and medical_report_en."
        )

    predictions = result.get("predictions") or []
    quality_metrics = result.get("quality_metrics") or {}
    heatmap_stats = result.get("heatmap_stats") or {}
    model_info = result.get("model_info") or {}

    try:
        primary_diagnosis = reports_en["prediction"]["primary_diagnosis"]
        primary_confidence = reports_en["prediction"]["confidence"]
    except (KeyError, TypeError) as exc:
        raise ValueError(
            "Invalid medical report prediction structure."
        ) from exc

    now = datetime.now(timezone.utc).isoformat()
    image_sha256 = hashlib.sha256(original_bytes).hexdigest()

    requested_language = str(result.get("language") or "ar").lower()
    selected_report = (
        reports_en if requested_language == "en" else reports_ar
    )

    row = {
        "id": analysis_id,
        "user_id": user.get("user_id") or None,
        "username": user.get("username") or None,
        "department": user.get("department") or None,
        "role": user.get("role") or None,
        "supervisor_id": user.get("supervisor_id") or None,
        "patient_name": patient.get("patient_name") or "",
        "birth_date": patient.get("birth_date") or "",
        "gender": patient.get("gender") or "",
        "age": patient.get("age"),
        "primary_diagnosis": primary_diagnosis,
        "primary_confidence": primary_confidence,
        "top_findings": predictions,
        "original_image_url": original_url,
        "heatmap_url": visual_urls.get("heatmap_enhanced"),
        "overlay_url": visual_urls.get("overlay"),
        "contour_url": visual_urls.get("contour"),
        "threshold_url": visual_urls.get("threshold"),
        "comparison_url": visual_urls.get("comparison"),
        "medical_report_json": json.dumps(
            {
                **selected_report,
                "report_integrity": result.get("report_integrity") or {},
                "janus": result.get("janus") or {},
                "safety": result.get("safety") or {},
            },
            ensure_ascii=False,
        ),
        "report_ar": reports_ar,
        "report_en": reports_en,
        "patient_data": _localized_patient(patient),
        "predictions_data": predictions,
        "quality_metrics": quality_metrics,
        "heatmap_stats": heatmap_stats,
        "model_info": model_info,
        "visualization_urls": visual_urls,
        "language_requested": requested_language,
        "original_filename": original_filename,
        "image_sha256": image_sha256,
        "server_version": model_info.get("version", "unknown"),
        "processing_time_ms": model_info.get("inference_time_ms"),
        "status": "completed",
        "created_at": now,
        "updated_at": now,
    }

    logger.info(
        "DB: inserting analyses row analysis_id=%s columns=%d",
        analysis_id,
        len(row),
    )

    # PostgREST schemas can differ slightly between the legacy Space and a
    # new Supabase project. Retry by removing only a column explicitly reported
    # as unknown, so a schema mismatch does not discard the already-uploaded
    # images and URLs.
    candidate = dict(row)
    last_error = None
    for _ in range(20):
        try:
            rows = _json_request(
                "POST",
                "analyses",
                candidate,
                params={"on_conflict": "id"},
            )
            saved_row = rows[0] if isinstance(rows, list) and rows else candidate
            logger.info("DB: SAVE SUCCESS analysis_id=%s", analysis_id)
            return {
                "saved": True,
                "analysis_id": analysis_id,
                "row": saved_row,
                "original_image_url": original_url,
                "visualization_urls": visual_urls,
            }
        except Exception as exc:
            last_error = exc
            message = str(exc)
            match = re.search(r'column ["\']([^"\']+)["\'] of relation ["\']analyses["\'] does not exist', message, re.I)
            if not match:
                break
            bad_column = match.group(1)
            if bad_column not in candidate:
                break
            logger.warning("DB: removing unknown analyses column=%s and retrying", bad_column)
            candidate.pop(bad_column, None)

    # Storage has already succeeded. Return the URLs even if the relational
    # insert failed; the API can still display the original X-ray and Grad-CAM
    # images and report the exact database error to the client.
    logger.exception("DB: relational save failed analysis_id=%s", analysis_id)
    return {
        "saved": False,
        "analysis_id": analysis_id,
        "row": None,
        "original_image_url": original_url,
        "visualization_urls": visual_urls,
        "error": str(last_error) if last_error else "unknown_database_error",
        "error_type": type(last_error).__name__ if last_error else "RuntimeError",
        "storage_saved": True,
    }


def health_check() -> Dict[str, Any]:
    """Verify REST access and Storage bucket availability without exposing secrets."""
    if not enabled():
        return {"ok": False, "reason": "database_not_configured", **config_info()}
    try:
        rows = _json_request("GET", "analyses", params={"select": "id", "limit": "1"})
        _ensure_bucket()
        return {"ok": True, "analyses_accessible": True, "bucket_accessible": True, **config_info()}
    except Exception as exc:
        logger.exception("DB HEALTH CHECK FAILED")
        return {
            "ok": False,
            "analyses_accessible": False,
            "bucket_accessible": False,
            "error": str(exc),
            "error_type": type(exc).__name__,
            **config_info(),
        }


def get_analysis(analysis_id: str):
    analysis_id = (analysis_id or "").strip()
    if not analysis_id:
        raise ValueError("analysis_id is required.")

    return_rows = _json_request(
        "GET",
        "analyses",
        params={
            "id": f"eq.{analysis_id}",
            "limit": "1",
        },
    )

    return return_rows[0] if return_rows else None


def list_analyses(user_id: str, limit: int = 50, offset: int = 0):
    user_id = (user_id or "").strip()
    if not user_id:
        raise ValueError("user_id is required.")

    safe_limit = min(max(int(limit), 1), 200)
    safe_offset = max(int(offset), 0)

    return _json_request(
        "GET",
        "analyses",
        params={
            "user_id": f"eq.{user_id}",
            "order": "created_at.desc",
            "limit": str(safe_limit),
            "offset": str(safe_offset),
        },
    )


def delete_analysis(analysis_id: str, user_id: str):
    analysis_id = (analysis_id or "").strip()
    user_id = (user_id or "").strip()

    if not analysis_id:
        raise ValueError("analysis_id is required.")
    if not user_id:
        raise ValueError("user_id is required.")

    logger.info(
        "DB: deleting analysis_id=%s user_id=%s",
        analysis_id,
        user_id,
    )

    try:
        response = requests.delete(
            _rest_url("analyses"),
            headers={
                **_headers(),
                "Prefer": "return=minimal",
            },
            params={
                "id": f"eq.{analysis_id}",
                "user_id": f"eq.{user_id}",
            },
            timeout=REQUEST_TIMEOUT,
        )
    except requests.RequestException as exc:
        logger.exception("DB: delete network error")
        raise RuntimeError(f"Delete request failed: {exc}") from exc

    if response.status_code >= 300:
        body = _safe_response_text(response, 1000)
        logger.error(
            "DB: delete failed status=%s body=%s",
            response.status_code,
            body,
        )
        raise RuntimeError(
            f"Delete failed ({response.status_code}): {body}"
        )

    logger.info("DB: delete OK analysis_id=%s", analysis_id)
    return True