"""Ingest the new structured ADC scrape into a running LightRAG server.

Input records are produced by scripts/adc_scraper/scrape_adc_product_data.py.
The script sends each record's `lightrag_text` to `/documents/text` and uses
the ADC detail URL as `file_source`, so repeated runs can detect duplicates.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
import urllib.error
import urllib.request
from collections import Counter
from pathlib import Path
from typing import Any


DEFAULT_INPUT = Path("data_adc_databaza/adc_scrape_2026_05_04/adc_products_structured.json")
DEFAULT_PROGRESS = Path("outputs/lightrag_ingest/adc_structured_progress.json")
DEFAULT_LIGHTRAG_URL = "http://localhost:9621"

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")


def load_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def write_json(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")


def load_progress(path: Path, resume: bool) -> dict[str, Any]:
    if resume and path.exists():
        return load_json(path)
    return {"done": [], "failed": [], "skipped": Counter(), "last_track_ids": []}


def save_progress(path: Path, progress: dict[str, Any]) -> None:
    payload = {
        "done": sorted(set(progress.get("done", []))),
        "failed": sorted(set(progress.get("failed", []))),
        "skipped": dict(progress.get("skipped", {})),
        "last_track_ids": progress.get("last_track_ids", [])[-50:],
    }
    write_json(path, payload)


def healthcheck(base_url: str) -> dict[str, Any]:
    with urllib.request.urlopen(f"{base_url}/health", timeout=10) as response:
        return json.loads(response.read().decode("utf-8"))


def pipeline_status(base_url: str) -> dict[str, Any]:
    with urllib.request.urlopen(f"{base_url}/documents/pipeline_status", timeout=10) as response:
        return json.loads(response.read().decode("utf-8"))


def post_text(base_url: str, text: str, file_source: str, timeout: int) -> dict[str, Any]:
    payload = json.dumps({"text": text, "file_source": file_source}, ensure_ascii=False).encode("utf-8")
    request = urllib.request.Request(
        f"{base_url}/documents/text",
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return json.loads(response.read().decode("utf-8"))


def product_name(record: dict[str, Any]) -> str:
    product = record.get("product") or {}
    graph_hints = record.get("graph_hints") or {}
    return str(graph_hints.get("drug") or product.get("name") or product.get("slug") or "unknown")


def source_id(record: dict[str, Any]) -> str:
    urls = record.get("urls") or {}
    product = record.get("product") or {}
    return str(urls.get("detail") or product.get("product_id") or product_name(record))


def is_clinical(record: dict[str, Any]) -> bool:
    hints = record.get("graph_hints") or {}
    return any(
        hints.get(key)
        for key in (
            "contraindications_text",
            "interactions_text",
            "warnings_text",
            "side_effects_text",
        )
    )


CLINICAL_KEYS = (
    "contraindications_text",
    "interactions_text",
    "warnings_text",
    "side_effects_text",
)


def clinical_text_length(record: dict[str, Any]) -> int:
    hints = record.get("graph_hints") or {}
    return sum(len(str(hints.get(key) or "")) for key in CLINICAL_KEYS)


def dedupe_key(record: dict[str, Any], mode: str) -> str:
    """Group key for records that describe the same medicine.

    `substance` groups products by their set of active substances (different
    packages and brands of the same substance share one leaflet content);
    products without known substances fall back to the clinical text itself.
    `text` groups only products whose clinical texts are identical.
    """
    hints = record.get("graph_hints") or {}
    if mode == "substance":
        substances = hints.get("active_substances") or []
        if isinstance(substances, str):
            substances = [substances]
        names = sorted({str(item).strip().lower() for item in substances if str(item).strip()})
        if names:
            return "substance:" + "|".join(names)
    if mode in {"substance", "text"}:
        texts = [str(hints.get(key) or "").strip() for key in CLINICAL_KEYS]
        return "text:" + json.dumps(texts, ensure_ascii=False)
    return "source:" + source_id(record)


def iter_selected_records(
    records: list[dict[str, Any]],
    progress: dict[str, Any],
    clinical_only: bool,
    dedupe: str = "none",
    retry_failed: bool = False,
) -> tuple[list[dict[str, Any]], Counter]:
    done = set(progress.get("done", []))
    failed = set() if retry_failed else set(progress.get("failed", []))
    skipped: Counter = Counter()
    candidates: list[dict[str, Any]] = []

    for record in records:
        if clinical_only and not is_clinical(record):
            skipped["not_clinical"] += 1
            continue
        if not str(record.get("lightrag_text") or "").strip():
            skipped["missing_lightrag_text"] += 1
            continue
        candidates.append(record)

    # Group duplicates; a group already covered by an uploaded record is skipped.
    groups: dict[str, list[dict[str, Any]]] = {}
    for record in candidates:
        groups.setdefault(dedupe_key(record, dedupe), []).append(record)

    selected: list[dict[str, Any]] = []
    for members in groups.values():
        sources = {source_id(record) for record in members}
        if sources & (done | failed):
            skipped["already_processed"] += len(members)
            continue
        # Representative: the record with the richest clinical text.
        selected.append(max(members, key=clinical_text_length))
        skipped["duplicate"] += len(members) - 1

    return selected, skipped


def show_status(base_url: str, progress_path: Path) -> None:
    if progress_path.exists():
        progress = load_json(progress_path)
        print(f"Progress file: {progress_path}")
        print(f"Done: {len(progress.get('done', []))}")
        print(f"Failed: {len(progress.get('failed', []))}")
        print(f"Skipped: {progress.get('skipped', {})}")
    else:
        print(f"No progress file found: {progress_path}")

    try:
        status = pipeline_status(base_url)
    except Exception as exc:
        print(f"LightRAG pipeline status unavailable: {exc}")
        return

    print(f"Pipeline busy: {status.get('busy')}")
    print(f"Latest message: {status.get('latest_message', '')}")


def ingest(args: argparse.Namespace) -> None:
    records = load_json(args.input)
    if not isinstance(records, list):
        raise SystemExit(f"Input must be a JSON list: {args.input}")

    progress = load_progress(args.progress, resume=args.resume)
    selected, skipped = iter_selected_records(
        records,
        progress,
        clinical_only=not args.all_records,
        dedupe=args.dedupe,
        retry_failed=args.retry_failed,
    )
    progress["skipped"] = Counter(progress.get("skipped", {})) + skipped

    if args.limit is not None:
        selected = selected[: args.limit]

    print(f"Input records: {len(records)}")
    print(f"Selected for upload: {len(selected)}")
    print(f"Skipped now: {dict(skipped)}")
    print(f"Progress: done={len(progress.get('done', []))}, failed={len(progress.get('failed', []))}")

    if args.dry_run:
        for record in selected[:10]:
            print(f"- {product_name(record)[:90]} | {source_id(record)}")
        save_progress(args.progress, progress)
        return

    health = healthcheck(args.url)
    configuration = health.get("configuration") or {}
    print(f"LightRAG OK: {args.url}")
    if configuration:
        print(f"LLM: {configuration.get('llm_model')} | Embedding: {configuration.get('embedding_model')}")

    done = set(progress.get("done", []))
    failed = set(progress.get("failed", []))
    last_track_ids = list(progress.get("last_track_ids", []))

    for index, record in enumerate(selected, start=1):
        source = source_id(record)
        text = str(record.get("lightrag_text") or "").strip()
        name = product_name(record)

        try:
            result = post_text(args.url, text=text, file_source=source, timeout=args.timeout)
            status = result.get("status", "unknown")
            track_id = result.get("track_id")
            if track_id:
                last_track_ids.append(track_id)
            if status in {"success", "duplicated"}:
                done.add(source)
                mark = "OK"
            else:
                failed.add(source)
                mark = "!!"
            print(f"[{index:5d}/{len(selected):5d}] {mark} {status:<10} {name[:80]}")
        except urllib.error.HTTPError as exc:
            failed.add(source)
            detail = exc.read().decode("utf-8", errors="replace")[:300]
            print(f"[{index:5d}/{len(selected):5d}] HTTP {exc.code} {name[:80]} | {detail}")
        except Exception as exc:
            failed.add(source)
            print(f"[{index:5d}/{len(selected):5d}] ERROR {name[:80]} | {exc}")

        progress["done"] = sorted(done)
        progress["failed"] = sorted(failed)
        progress["last_track_ids"] = last_track_ids[-50:]

        if index % args.save_every == 0:
            save_progress(args.progress, progress)
            print(f"Saved progress: done={len(done)}, failed={len(failed)}")

        time.sleep(args.delay)

    save_progress(args.progress, progress)
    print(f"Finished upload loop. Done={len(done)}, failed={len(failed)}")
    print("LightRAG will continue extracting graph entities and relations in the background.")


def main() -> None:
    parser = argparse.ArgumentParser(description="Ingest structured ADC data into LightRAG.")
    parser.add_argument("--input", type=Path, default=DEFAULT_INPUT)
    parser.add_argument("--progress", type=Path, default=DEFAULT_PROGRESS)
    parser.add_argument("--url", default=DEFAULT_LIGHTRAG_URL)
    parser.add_argument("--limit", type=int, default=50)
    parser.add_argument("--delay", type=float, default=0.3)
    parser.add_argument("--timeout", type=int, default=60)
    parser.add_argument("--save-every", type=int, default=10)
    parser.add_argument("--resume", action="store_true", help="Resume from the progress file.")
    parser.add_argument("--all-records", action="store_true", help="Upload all records, not only clinical ones.")
    parser.add_argument(
        "--dedupe",
        choices=("none", "text", "substance"),
        default="substance",
        help="Upload one representative per active substance set (default) or per identical clinical text.",
    )
    parser.add_argument("--retry-failed", action="store_true", help="Upload again records marked as failed.")
    parser.add_argument("--dry-run", action="store_true", help="Show selected records without uploading.")
    parser.add_argument("--status", action="store_true", help="Show local progress and LightRAG pipeline status.")
    args = parser.parse_args()

    if args.status:
        show_status(args.url, args.progress)
        return

    ingest(args)


if __name__ == "__main__":
    main()