"""BEAM-10M dataset loader — fetches real conversations from HF.

BEAM-10M (Mohammadta/BEAM-10M) is a LongMemEval-style benchmark:
each row is a long, multi-session conversation with:
  - conversation_id
  - user_profile (user_info + user_relationships) — ground-truth facts
  - narratives — a structured timeline of events
  - chat — 10 plans, each with 10 batches of ~60 turns
  - probing_questions — questions to test memory of the conversation
  - plans — alternate plans structure

The HuggingFace datasets API is rate-limited from many sandboxes
(429 Too Many Requests from CloudFront). The datasets-server endpoint
(hosted separately) IS reachable. This loader pulls via datasets-server
so the bench works even when direct HF datasets API is blocked.

BULK DOWNLOAD PATHS (added for full-dataset benchmarks):
  - Local cache directory of per-row JSON files (preferred — produced
    by scripts/download_beam_full.sh OR by the .github/workflows/
    beam-cache.yml workflow that runs on a GitHub Actions runner)
  - Local parquet file (if you have downloaded the full parquet via
    `huggingface-cli download Mohammadta/BEAM-10M` on a runner whose
    IP is not rate-limited; the file is ~342MB). Pass the path via
    the BEAM_PARQUET env var or the parquet_path argument.
  - datasets-server /rows endpoint as fallback (works in sandboxes)

Usage:
    from cortexm.bench.beam_loader import load_beam_rows
    rows = load_beam_rows(n=10)  # fetch all 10 conversations
    for r in rows:
        print(r['conversation_id'])
        print(r['user_profile']['user_info'])
"""
from __future__ import annotations

import json
import os
import urllib.request
from pathlib import Path
from typing import Iterator


DATASETS_SERVER = "https://datasets-server.huggingface.co"
DATASET_NAME = "Mohammadta/BEAM-10M"
CONFIG = "default"
SPLIT = "10M"
TOTAL_ROWS = 10  # the 10M split has 10 conversations, each ~10M tokens


def _fetch_rows(offset: int = 0, length: int = 1) -> dict:
    """Fetch BEAM-10M rows via the datasets-server endpoint."""
    url = (f"{DATASETS_SERVER}/rows?dataset={DATASET_NAME}"
           f"&config={CONFIG}&split={SPLIT}"
           f"&offset={offset}&length={length}")
    req = urllib.request.Request(
        url, headers={"User-Agent": "context-m-bench/1.0"})
    with urllib.request.urlopen(req, timeout=600) as resp:
        return json.loads(resp.read().decode("utf-8"))


def _load_parquet(parquet_path: str | Path, n: int = TOTAL_ROWS) -> list[dict]:
    """Load rows from a local parquet file (requires pyarrow or pandas).

    Each parquet row has the same shape as the datasets-server /rows
    endpoint's `row` field. We return just the row dicts (no envelope).
    """
    try:
        import pyarrow.parquet as pq
    except ImportError:
        try:
            import pandas as pd
        except ImportError:
            raise ImportError(
                "loading BEAM-10M from parquet requires pyarrow or "
                "pandas — install with `pip install pyarrow` or "
                "`pip install pandas`")
        df = pd.read_parquet(parquet_path)
        rows = []
        for _, r in df.head(n).iterrows():
            rows.append(r.to_dict() if hasattr(r, "to_dict") else dict(r))
        return rows
    table = pq.read_table(parquet_path)
    n = min(n, table.num_rows)
    rows = []
    # batch-convert to Python dicts (column-at-a-time is faster than row)
    cols = {name: table.column(name).to_pylist() for name in table.column_names}
    for i in range(n):
        rows.append({name: cols[name][i] for name in cols})
    return rows


def load_beam_rows(n: int = 2, *, cache_dir: str | None = None,
                   parquet_path: str | None = None) -> list[dict]:
    """Load N BEAM-10M conversations.

    Resolution order (first available wins):
      1. parquet_path argument (or BEAM_PARQUET env var) — a single
         parquet file containing all 10 rows
      2. cache_dir/beam_row_<i>.json — per-row JSON files (the format
         produced by download_beam_full.sh and beam-cache.yml)
      3. datasets-server /rows endpoint — streamed on demand

    Each row is a dict with keys: conversation_id, conversation_seed,
    narratives, user_profile, conversation_plan, user_questions, chat,
    probing_questions, plans.

    The full 10M dataset has 10 conversations totaling ~975MB in
    memory. We fetch one row at a time (each is ~50-110MB).
    """
    # 1. parquet path?
    pq_path = parquet_path or os.environ.get("BEAM_PARQUET")
    if pq_path and Path(pq_path).exists():
        return _load_parquet(pq_path, n=min(n, TOTAL_ROWS))

    cache = Path(cache_dir) if cache_dir else None
    if cache:
        cache.mkdir(parents=True, exist_ok=True)

    rows: list[dict] = []
    n = min(n, TOTAL_ROWS)
    for i in range(n):
        # 2. cached row?
        cached_path = cache / f"beam_row_{i}.json" if cache else None
        if cached_path and cached_path.exists():
            try:
                data = json.loads(cached_path.read_text())
                if "rows" in data and data["rows"]:
                    rows.append(data["rows"][0]["row"])
                    continue
                # some manifests store the row directly without envelope
                if "conversation_id" in data:
                    rows.append(data)
                    continue
            except (json.JSONDecodeError, KeyError):
                # corrupt cache — re-fetch
                cached_path.unlink(missing_ok=True)
        # 3. fetch from datasets-server
        data = _fetch_rows(offset=i, length=1)
        if cached_path:
            cached_path.write_text(json.dumps(data))
        if "rows" in data and data["rows"]:
            rows.append(data["rows"][0]["row"])
    return rows


def parse_user_facts(row: dict) -> list[dict]:
    """Extract ground-truth facts from a BEAM row's user_profile.

    The user_profile has:
      user_info: free text "Name: Jennifer Mccall / Age: 59 / ..."
      user_relationships: free text with bullet points

    We extract structured facts via regex (zero LLM calls).
    """
    import re
    up = row.get("user_profile", {})
    info = up.get("user_info", "") or ""
    rels = up.get("user_relationships", "") or ""

    facts: list[dict] = []
    conv_id = row.get("conversation_id", "unknown")
    user_id = f"beam_{conv_id}"

    # name
    m = re.search(r"Name:\s*([^\n]+)", info)
    if m:
        facts.append({"subject": user_id, "relation": "name",
                       "value": m.group(1).strip()})
    # age
    m = re.search(r"Age:\s*(\d+)", info)
    if m:
        facts.append({"subject": user_id, "relation": "age",
                       "value": m.group(1).strip()})
    # gender
    m = re.search(r"Gender:\s*([^\n]+)", info)
    if m:
        facts.append({"subject": user_id, "relation": "gender",
                       "value": m.group(1).strip()})
    # location
    m = re.search(r"Location:\s*([^\n]+)", info)
    if m:
        facts.append({"subject": user_id, "relation": "location",
                       "value": m.group(1).strip()})
    # profession
    m = re.search(r"Profession:\s*([^\n]+)", info)
    if m:
        facts.append({"subject": user_id, "relation": "profession",
                       "value": m.group(1).strip()})

    # relationships — extract names + relations
    # relationship sections start with all-caps headers like
    # "PARENTS & GUARDIANS:" "ROMANTIC PARTNER:" "CHILDREN:" etc.
    cur_section = None
    for line in rels.split("\n"):
        line = line.strip()
        if not line:
            continue
        # detect header — line is all caps + colon
        if line.endswith(":") and line[:-1] == line[:-1].upper():
            cur_section = line[:-1]
            continue
        # bullet line — extract name + age
        m = re.match(r"•\s*([^\(]+)\s*\(.*?\)", line)
        if m and cur_section:
            name = m.group(1).strip().rstrip(",")
            # convert section to relation
            rel_map = {
                "PARENTS & GUARDIANS": "parent",
                "ROMANTIC PARTNER": "partner",
                "CHILDREN": "child",
                "SIBLINGS": "sibling",
                "FRIENDS": "friend",
                "COLLEAGUES": "colleague",
            }
            relation = rel_map.get(cur_section, cur_section.lower())
            facts.append({"subject": user_id, "relation": relation,
                           "value": name})
    return facts


def parse_chat_turns(row: dict) -> list[dict]:
    """Flatten a row's chat history into a list of user turns.

    Each turn has: content, id, index, question_type, role, time_anchor.
    We only keep role='user' turns (these are the messages the user
    actually said — what Context-M needs to extract facts from).
    """
    chat = row.get("chat", [])
    if not isinstance(chat, list):
        return []
    turns_out: list[dict] = []
    for plan_obj in chat:
        if not isinstance(plan_obj, dict):
            continue
        # plan_obj has one key like "plan-1"
        for plan_name, batches in plan_obj.items():
            if not isinstance(batches, list):
                continue
            for batch in batches:
                if not isinstance(batch, dict):
                    continue
                turns = batch.get("turns", [])
                if not isinstance(turns, list):
                    continue
                for turn_list in turns:
                    if not isinstance(turn_list, list):
                        continue
                    for turn in turn_list:
                        if not isinstance(turn, dict):
                            continue
                        if turn.get("role") == "user":
                            turns_out.append({
                                "content": turn.get("content", ""),
                                "time_anchor": turn.get("time_anchor"),
                                "question_type": turn.get("question_type"),
                                "plan": plan_name,
                                "batch": batch.get("batch_number"),
                            })
    return turns_out


def beam_rows_to_personas(rows: list[dict], *,
                           max_turns_per_persona: int = 100,
                           include_profile: bool = True) -> list[dict]:
    """Convert BEAM rows into the persona dict format the bench expects.

    Returns list of: {user_id, text, facts} where:
      user_id — derived from conversation_id
      text — concatenation of user_profile (if include_profile=True)
             + the first N user turns. The user_profile contains the
             ground-truth facts (Name, Age, Location, etc.) which the
             μ=0 extractor should pull from the explicit "Name: X"
             lines; the chat turns are the conversational context.
      facts — ground-truth structured facts parsed from user_profile
              (separate from text so the bench can check recall)
    """
    personas = []
    for row in rows:
        conv_id = row.get("conversation_id", "unknown")
        user_id = f"beam_{conv_id}"
        facts = parse_user_facts(row)
        turns = parse_chat_turns(row)
        # cap turns to keep ingest tractable
        turns = turns[:max_turns_per_persona]
        # build the ingest text: profile first (so facts are stated
        # explicitly), then the conversation turns
        parts = []
        if include_profile:
            up = row.get("user_profile", {})
            info = up.get("user_info", "")
            rels = up.get("user_relationships", "")
            if info:
                parts.append(info)
            if rels:
                parts.append(rels)
        # append the chat turns
        for t in turns:
            if t.get("content"):
                parts.append(t["content"])
        text = "\n".join(parts)
        personas.append({
            "user_id": user_id,
            "text": text,
            "facts": facts,
            "n_turns": len(turns),
            "conversation_id": conv_id,
        })
    return personas


__all__ = [
    "load_beam_rows",
    "parse_user_facts",
    "parse_chat_turns",
    "beam_rows_to_personas",
    "DATASET_NAME",
]