"""
NeuroJenML — Automated paper pre-screening (free Gemma).

Before a paper reaches the human review queue, an automated screen gates it on
three checks (see DATA_HANDLING.md §3.1):

  1. has_content   — enough extractable text to be a real paper, not a stub/scan.
  2. is_ad_related — about Alzheimer's disease / neurodegeneration.
  3. is_scientific — a scientific study (methods/results), not an op-ed/news/ad.

A paper passes only if ALL THREE are true. Passing papers move to `pending`
(human queue); failing papers move to `screened_out` with a reason, so a human
never has to triage obvious junk.

Backend: free Gemma via the HF router (no MedGemma needed for this cheap gate).
Degrades gracefully — if HF_TOKEN is missing or the call fails, it falls back to
a conservative heuristic screen so the pipeline never stalls.
"""

import os
import json

import httpx

# Minimum characters of real text to be considered a paper (not a scanned image
# or a near-empty stub).
MIN_CONTENT_CHARS = 600

# Screening uses TWO models in parallel to reduce hallucination risk on the
# three-check gate. Both must agree for the paper to pass; on disagreement the
# result is conservative (fail unless the difference is immaterial — see
# _merge_verdicts below). Running two small models is cheaper and more reliable
# than trusting one alone.
#
# Primary:   google/gemma-2-2b-it  — confirmed free-tier, user-verified no errors.
# Secondary: google/gemma-3-4b-it  — stronger, used as cross-check.
#
# Override the models:
#   SCREENING_MODEL_PRIMARY   (default: google/gemma-2-2b-it)
#   SCREENING_MODEL_SECONDARY (default: google/gemma-3-4b-it)
# Set both to the same model to disable dual-model mode.
_DEFAULT_PRIMARY_MODEL   = "google/gemma-2-2b-it"
_DEFAULT_SECONDARY_MODEL = "google/gemma-3-4b-it"

# Legacy single-model env var still honoured — sets the primary if present.
_DEFAULT_SCREENING_MODEL = _DEFAULT_PRIMARY_MODEL

# AD / neurodegeneration signal terms for the heuristic fallback.
_AD_TERMS = [
    "alzheimer", "dementia", "amyloid", "tau", "neurodegener", "cognitive decline",
    "neurofibrillary", "apoe", "mci", "mild cognitive impairment", "beta-amyloid",
    "abeta", "a\u03b2", "neuroinflammation", "tauopathy", "hippocamp",
]

# Markers that a document is a scientific study rather than prose/news.
_SCIENCE_TERMS = [
    "method", "results", "conclusion", "abstract", "hypothesis", "p <", "p<",
    "p =", "cohort", "sample", "statistical", "experiment", "assay", "control group",
    "figure", "table", "doi", "et al", "regression", "significant",
]

_SCREEN_PROMPT = (
    "You are a screening gate for an Alzheimer's disease (AD) research dataset. "
    "Read the document excerpt and decide three booleans, then return STRICT JSON only:\n"
    '  "has_content": true if it contains substantive readable text (not an empty/'
    'corrupted/scanned-image stub).\n'
    '  "is_ad_related": true if it concerns Alzheimer\'s disease, dementia, or '
    "neurodegeneration (central or via a systemic/peripheral link).\n"
    '  "is_scientific": true if it is a scientific study or review (has methods, '
    "data, or structured findings) rather than news, opinion, marketing, or fiction.\n"
    '  "reason": one short sentence explaining the decision.\n'
    "Return ONLY the JSON object."
)


def _heuristic_screen(text: str) -> dict:
    """Conservative keyword-based screen used when the LLM is unavailable."""
    lower = (text or "").lower()
    has_content = len((text or "").strip()) >= MIN_CONTENT_CHARS and not lower.startswith("[pdf binary")
    is_ad_related = any(term in lower for term in _AD_TERMS)
    science_hits = sum(1 for term in _SCIENCE_TERMS if term in lower)
    is_scientific = science_hits >= 3
    passed = has_content and is_ad_related and is_scientific
    return {
        "passed": passed,
        "checks": {
            "has_content": has_content,
            "is_ad_related": is_ad_related,
            "is_scientific": is_scientific,
        },
        "reason": _reason(has_content, is_ad_related, is_scientific),
        "model": "heuristic",
    }


def _reason(has_content: bool, ad: bool, sci: bool) -> str:
    if not has_content:
        return "Insufficient extractable text — likely an empty, stub, or scanned document."
    if not ad:
        return "Does not appear to be related to Alzheimer's disease or neurodegeneration."
    if not sci:
        return "Does not appear to be a scientific study (missing methods/data/structure)."
    return "Passed: substantive, AD-related, scientific content."


def _coerce_verdict(obj: dict, model: str) -> dict:
    checks = {
        "has_content": bool(obj.get("has_content")),
        "is_ad_related": bool(obj.get("is_ad_related")),
        "is_scientific": bool(obj.get("is_scientific")),
    }
    passed = all(checks.values())
    reason = obj.get("reason") or _reason(**{
        "has_content": checks["has_content"],
        "ad": checks["is_ad_related"],
        "sci": checks["is_scientific"],
    })
    return {"passed": passed, "checks": checks, "reason": str(reason)[:300], "model": model}


def _parse_json_lenient(raw: str) -> dict:
    if not raw:
        return {}
    try:
        return json.loads(raw)
    except Exception:
        pass
    start = raw.find("{")
    if start == -1:
        return {}
    depth = 0
    for i in range(start, len(raw)):
        if raw[i] == "{":
            depth += 1
        elif raw[i] == "}":
            depth -= 1
            if depth == 0:
                try:
                    return json.loads(raw[start:i + 1])
                except Exception:
                    return {}
    return {}


async def _call_model(token: str, model: str, prompt: str) -> dict:
    """Call one HF model and return a coerced verdict, or None on failure."""
    try:
        from huggingface_hub import AsyncInferenceClient
        client = AsyncInferenceClient(api_key=token)
        resp = await client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.0,
            max_tokens=300,
        )
        raw = resp.choices[0].message.content
        parsed = _parse_json_lenient(raw)
        return _coerce_verdict(parsed, model) if parsed else None
    except Exception:
        return None


def _merge_verdicts(a: dict, b: dict) -> dict:
    """Consensus logic: conservative — FAIL if any model raises a check failure.

    Rationale: false negatives (passing bad papers) are worse than false
    positives (rejecting good ones). A human can always override a screened-out
    paper; a hallucinated pass can corrupt the dataset.

    Both models must agree on PASS for the paper to pass.  If only one model
    returned a result, its verdict is used directly.
    """
    if a is None and b is None:
        return None  # both failed → fall through to heuristic
    if a is None:
        return b
    if b is None:
        return a

    # Each check: True only if both models say True.
    merged_checks = {
        k: a["checks"].get(k, False) and b["checks"].get(k, False)
        for k in ("has_content", "is_ad_related", "is_scientific")
    }
    passed = all(merged_checks.values())
    # Prefer the more informative reason (the one that failed, if any).
    reason = (
        b["reason"] if not merged_checks.get("is_ad_related") or not merged_checks.get("is_scientific")
        else a["reason"]
    )
    return {
        "passed": passed,
        "checks": merged_checks,
        "reason": str(reason)[:300],
        "model": f"{a['model']}+{b['model']}",
    }


async def screen(text: str) -> dict:
    """Screen a paper. Returns {passed, checks, reason, model}.

    Runs TWO models in parallel and merges via conservative consensus to reduce
    hallucination risk. Cheap length guard runs first; heuristic fallback fires
    if both model calls fail so the pipeline never stalls.
    """
    import asyncio as _asyncio

    text = text or ""
    # Hard fail on no content without spending any model call.
    if len(text.strip()) < MIN_CONTENT_CHARS or text.strip().lower().startswith("[pdf binary"):
        return {
            "passed": False,
            "checks": {"has_content": False, "is_ad_related": False, "is_scientific": False},
            "reason": "Insufficient extractable text — likely an empty, stub, or scanned document.",
            "model": "guard",
        }

    token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
    if not token:
        return _heuristic_screen(text)

    # Resolve model IDs. Legacy SCREENING_MODEL_ID overrides the primary.
    primary   = os.getenv("SCREENING_MODEL_ID") or os.getenv("SCREENING_MODEL_PRIMARY",   _DEFAULT_PRIMARY_MODEL)
    secondary = os.getenv("SCREENING_MODEL_SECONDARY", _DEFAULT_SECONDARY_MODEL)

    prompt = f"{_SCREEN_PROMPT}\n\n--- DOCUMENT EXCERPT ---\n{text[:6000]}\n--- END ---\n\nJSON:"

    async def _noop():
        return None

    # Fire both models concurrently.
    a, b = await _asyncio.gather(
        _call_model(token, primary, prompt),
        _call_model(token, secondary, prompt) if secondary != primary else _noop(),
        return_exceptions=False,
    )

    merged = _merge_verdicts(a, b)
    if merged is None:
        return _heuristic_screen(text)
    return merged