src/analyzer/data_loader.py

"""
data_loader.py — loads current grant snapshots (JSON) and optional past winners

This module keeps IO concerns simple and robust:
- Recursively loads current-grant JSON files under a snapshots directory
- Optionally loads past winners from either an Excel file or a JSON directory
- Returns Python lists of dictionaries; no model code here

Public API
---------
load_current_grants(snapshots_dir: Path | str, limit: int | None = None) -> list[dict]
load_past_winners(history_xlsx: Path | str | None = None,
                  history_json_dir: Path | str | None = None) -> list[dict]
"""
from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, List, Optional
import json
import logging

import pandas as pd

logger = logging.getLogger(__name__)

# ----------------------------- Current grants ---------------------------------

def load_current_grants(snapshots_dir: Path | str, limit: Optional[int] = None) -> List[Dict[str, Any]]:
    """Load current grant JSON snapshots from a directory tree.

    Each file is expected to be one JSON object. The function tolerates
    missing keys and will attach an `id` from the filename if not present.
    """
    snapshots_dir = Path(snapshots_dir)
    if not snapshots_dir.exists():
        logger.warning("snapshots directory not found: %s", snapshots_dir)
        return []

    records: List[Dict[str, Any]] = []
    for p in sorted(snapshots_dir.rglob("*.json")):
        try:
            with open(p, "r", encoding="utf-8") as f:
                rec = json.load(f)
            if not isinstance(rec, dict):
                logger.debug("Skipping non-object JSON: %s", p)
                continue
            rec.setdefault("id", p.stem)
            rec.setdefault("_path", str(p))
            records.append(rec)
            if limit and len(records) >= limit:
                break
        except Exception as e:  # pragma: no cover
            logger.warning("Failed to load %s: %s", p, e)
            continue

    logger.info("Loaded %d current grants from %s", len(records), snapshots_dir)
    return records

# ------------------------------ Past winners ----------------------------------

_CANON_COLS = {
    # canonical : candidate column names (lower/underscore)
    "project_title": ["project_title", "title", "name"],
    "abstract": ["abstract", "description", "summary", "public_description"],
    "competition": ["competition", "programme", "program"],
    "award_amount": ["award_amount", "amount", "grant", "project_cost", "award"],
    "lead_org": ["lead_org", "lead_organisation", "lead_organization", "organisation_name", "organization_name"],
    "year": ["year", "fy", "start_year"],
    "project_url": ["project_url", "url", "link"],
}


def _norm_cols(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df.columns = [
        (c if isinstance(c, str) else str(c))
        .lower()
        .replace(" ", "_")
        .replace("-", "_")
        for c in df.columns
    ]
    return df


def _ensure_canonical(df: pd.DataFrame) -> pd.DataFrame:
    for canon, candidates in _CANON_COLS.items():
        if canon in df.columns:
            continue
        for c in candidates:
            if c in df.columns:
                df[canon] = df[c]
                break
        if canon not in df.columns:
            df[canon] = None
    return df


def _load_past_winners_from_excel(xlsx_path: Path) -> List[Dict[str, Any]]:
    df = pd.read_excel(xlsx_path)
    df = _norm_cols(df)
    df = _ensure_canonical(df)
    return [row._asdict() if hasattr(row, "_asdict") else row.to_dict() for _, row in df.iterrows()]


def _load_past_winners_from_json_dir(json_dir: Path) -> List[Dict[str, Any]]:
    records: List[Dict[str, Any]] = []
    for p in sorted(json_dir.rglob("*.json")):
        try:
            with open(p, "r", encoding="utf-8") as f:
                rec = json.load(f)
            if not isinstance(rec, dict):
                continue
            rec.setdefault("_path", str(p))
            records.append(rec)
        except Exception:
            continue
    return records


def _load_past_winners_from_jsonl(jsonl_path: Path) -> List[Dict[str, Any]]:
    """Load past winners from JSONL file (optionally gzipped)."""
    import gzip

    records: List[Dict[str, Any]] = []

    # Check if file is gzipped
    open_func = gzip.open if str(jsonl_path).endswith('.gz') else open

    try:
        with open_func(jsonl_path, 'rt', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    rec = json.loads(line)
                    if isinstance(rec, dict):
                        records.append(rec)
                except json.JSONDecodeError:
                    continue
        return records
    except Exception as e:
        logger.warning(f"Failed to load JSONL from {jsonl_path}: {e}")
        return []


def load_past_winners(
    history_xlsx: Path | str | None = None,
    history_json_dir: Path | str | None = None,
    history_jsonl: Path | str | None = None
) -> List[Dict[str, Any]]:
    """Load past winners from Excel, JSONL, or JSON directory.

    Priority order: JSONL > JSON dir > Excel.
    If neither exists, returns an empty list.
    """
    # Try JSONL first (most efficient for large datasets)
    if history_jsonl is not None:
        jsonl = Path(history_jsonl)
        if jsonl.exists():
            recs = _load_past_winners_from_jsonl(jsonl)
            if recs:
                logger.info("Loaded %d past winners from JSONL: %s", len(recs), jsonl)
                return recs

    # Also check for past_winners.jsonl.gz in default location (for HF deployment)
    default_jsonl = Path("data/past_winners.jsonl.gz")
    if default_jsonl.exists() and history_jsonl is None:
        recs = _load_past_winners_from_jsonl(default_jsonl)
        if recs:
            logger.info("Loaded %d past winners from default JSONL: %s", len(recs), default_jsonl)
            return recs

    # JSON dir next
    if history_json_dir is not None:
        jdir = Path(history_json_dir)
        if jdir.exists():
            recs = _load_past_winners_from_json_dir(jdir)
            if recs:
                logger.info("Loaded %d past winners from JSON dir: %s", len(recs), jdir)
                return recs
            else:
                logger.info("No JSON past winners found under %s", jdir)

    # Excel last
    if history_xlsx is not None:
        xlsx = Path(history_xlsx)
        if xlsx.exists():
            recs = _load_past_winners_from_excel(xlsx)
            logger.info("Loaded %d past winners from Excel: %s", len(recs), xlsx)
            return recs
        else:
            logger.info("History Excel not found: %s", xlsx)

    return []


# Self-test
if __name__ == "__main__":
    import argparse
    logging.basicConfig(level=logging.INFO)
    parser = argparse.ArgumentParser()
    parser.add_argument("--snapshots-dir", type=Path, default=Path("data/snapshots"))
    parser.add_argument("--history-xlsx", type=Path, default=Path("data/past_winners.xlsx"))
    parser.add_argument("--limit", type=int, default=3)
    args = parser.parse_args()

    current = load_current_grants(args.snapshots_dir, limit=args.limit)
    history = load_past_winners(args.history_xlsx)
    print(f"current: {len(current)} | history: {len(history)}")
    if current:
        print("example current keys:", sorted(current[0].keys())[:12])
    if history:
        print("example history keys:", sorted(history[0].keys())[:12])