src/parsing.py

from __future__ import annotations

import json
import re
from typing import Any


_JSON_FENCE_PATTERN = re.compile(r"```json\s*(.*?)\s*```", re.IGNORECASE | re.DOTALL)
_CONSOLIDATED_FINDINGS_MARKER = "Consolidated list of findings:"


def _loads_first_json(candidate: str) -> Any:
    decoder = json.JSONDecoder()
    stripped = candidate.strip()
    if not stripped:
        raise ValueError("JSON candidate is empty.")

    for index, char in enumerate(stripped):
        if char not in "{[":
            continue
        try:
            payload, _ = decoder.raw_decode(stripped, index)
            return payload
        except json.JSONDecodeError:
            continue

    raise ValueError("Could not parse a JSON value from the candidate text.")


def _scan_json_values(text: str) -> list[Any]:
    decoder = json.JSONDecoder()
    values: list[Any] = []

    for index, char in enumerate(text):
        if char not in "{[":
            continue
        try:
            payload, _ = decoder.raw_decode(text, index)
        except json.JSONDecodeError:
            continue
        values.append(payload)

    return values


def _is_finding_item(value: Any) -> bool:
    if not isinstance(value, dict):
        return False
    return all(key in value for key in ("finding", "anatomical_location", "certainty"))


def _is_finding_list(value: Any) -> bool:
    if not isinstance(value, list):
        return False
    return all(_is_finding_item(item) for item in value)


def _normalize_discovery_candidate(payload: Any) -> dict[str, Any] | None:
    if isinstance(payload, dict) and "findings" in payload:
        if isinstance(payload["findings"], list):
            return payload
        return None

    if _is_finding_list(payload):
        return {"findings": payload}

    if _is_finding_item(payload):
        return None

    return None


def _has_structured_discovery_payload(values: list[Any]) -> bool:
    for value in values:
        if isinstance(value, dict) and "findings" in value:
            return True
        if _is_finding_list(value):
            return True
    return False


def strip_medgemma_thinking_trace(response: str) -> str:
    """Remove the notebook-documented MedGemma thinking trace when present."""
    if "<unused95>" in response:
        return response.split("<unused95>", 1)[1].lstrip()
    return response.strip()


def extract_finding_discovery_payload(response: str) -> dict[str, Any]:
    cleaned = strip_medgemma_thinking_trace(response).strip()

    scanned_values = _scan_json_values(cleaned)
    standalone_finding_dicts = [
        value for value in scanned_values if _is_finding_item(value)
    ]
    if len(standalone_finding_dicts) >= 2 and not _has_structured_discovery_payload(
        scanned_values
    ):
        raise ValueError(
            "Could not extract finding discovery JSON from MedGemma response. "
            f"Response preview: {cleaned[:500]!r}"
        )

    try:
        full_payload = json.loads(cleaned)
    except json.JSONDecodeError:
        full_payload = None
    else:
        normalized = _normalize_discovery_candidate(full_payload)
        if normalized is not None:
            return normalized

    for value in scanned_values:
        if isinstance(value, dict) and "findings" in value:
            normalized = _normalize_discovery_candidate(value)
            if normalized is not None:
                return normalized

    if _CONSOLIDATED_FINDINGS_MARKER in cleaned:
        after_marker = cleaned.split(_CONSOLIDATED_FINDINGS_MARKER, 1)[1].strip()
        try:
            marker_payload = _loads_first_json(after_marker)
        except ValueError:
            marker_payload = None
        else:
            normalized = _normalize_discovery_candidate(marker_payload)
            if normalized is not None:
                return normalized

    for match in _JSON_FENCE_PATTERN.finditer(cleaned):
        fenced_text = match.group(1).strip()
        try:
            fenced_payload = _loads_first_json(fenced_text)
        except ValueError:
            continue
        normalized = _normalize_discovery_candidate(fenced_payload)
        if normalized is not None:
            return normalized

    for value in scanned_values:
        if _is_finding_list(value):
            return {"findings": value}

    raise ValueError(
        "Could not extract finding discovery JSON from MedGemma response. "
        f"Response preview: {cleaned[:500]!r}"
    )


def extract_json_payload(response: str, expected_type: type) -> Any:
    cleaned = strip_medgemma_thinking_trace(response)

    candidates: list[str] = [
        match.group(1).strip()
        for match in _JSON_FENCE_PATTERN.finditer(cleaned)
    ]

    if not candidates:
        if expected_type is dict:
            start = cleaned.find("{")
            end = cleaned.rfind("}")
        elif expected_type is list:
            start = cleaned.find("[")
            end = cleaned.rfind("]")
        else:
            raise TypeError("expected_type must be dict or list.")

        if start == -1 or end == -1 or end <= start:
            raise ValueError("Could not locate a JSON payload in the model response.")
        candidates.append(cleaned[start : end + 1])

    errors: list[str] = []
    for candidate in candidates:
        try:
            payload = _loads_first_json(candidate)
        except ValueError as error:
            errors.append(str(error))
            continue

        if isinstance(payload, expected_type):
            return payload

        errors.append(
            f"expected {expected_type.__name__}, received {type(payload).__name__}"
        )

    detail = "; ".join(errors) if errors else "no JSON candidates were found"
    raise ValueError(
        f"Could not extract JSON {expected_type.__name__} from the model response: {detail}"
    )