"""US earnings-call transcripts from the kurry/sp500_earnings_transcripts dataset.

33,362 transcripts across 685 S&P 500 companies, 2005-2025, MIT-licensed. This is the
US counterpart to the BSE concall PDFs on the Indian side — except the text is
already speaker-tagged, so management commentary and analyst Q&A are separable
without parsing a PDF.

The parquet is ~1.8 GB (two row groups, the first holding 32,768 rows), so it is never
loaded whole. Listings read only the five metadata columns (milliseconds). The text of
one company is extracted by streaming the file in batches and keeping that symbol's
rows, which holds a few hundred rows in memory instead of the file.
"""
from __future__ import annotations

import functools
import re

REPO = "kurry/sp500_earnings_transcripts"
FILE = "parquet_files/part-0.parquet"
COVERAGE = (f"{REPO} (MIT): about 33,000 earnings calls for 685 S&P 500 companies, "
            "2005-2025; year and quarter are the dataset's own labels")

_META_COLS = ["symbol", "quarter", "year", "date", "company_name"]
_path_cache: str | None = None


def _path() -> str | None:
    """Local path of the parquet, downloading on first use (~1.8 GB, cached by HF).

    A failed download is not remembered, so a transient network error does not
    disable the tool until restart."""
    global _path_cache
    if _path_cache:
        return _path_cache
    try:
        from huggingface_hub import hf_hub_download
        _path_cache = hf_hub_download(REPO, FILE, repo_type="dataset")
    except Exception:
        return None
    return _path_cache


def available() -> bool:
    return _path() is not None


@functools.lru_cache(maxsize=1)
def _meta_cached(path: str):
    import pyarrow.parquet as pq
    df = pq.read_table(path, columns=_META_COLS).to_pandas()
    df["date"] = df["date"].astype(str).str[:10]
    return df


def _meta():
    p = _path()
    return _meta_cached(p) if p else None


def symbols() -> frozenset:
    m = _meta()
    return frozenset(m["symbol"].dropna()) if m is not None else frozenset()


def listing(symbol: str, year: int | None = None, limit: int = 100) -> list[dict]:
    """Calls available for a symbol, newest first — metadata only, not the text."""
    m = _meta()
    if m is None:
        return []
    df = m[m["symbol"] == symbol.upper()]
    if year is not None:
        df = df[df["year"] == int(year)]
    df = df.sort_values(["date", "year", "quarter"], ascending=False).head(limit)
    return [{"company": r.company_name, "year": int(r.year), "quarter": int(r.quarter),
             "date": r.date} for r in df.itertuples()]


@functools.lru_cache(maxsize=8)
def _calls(path: str, symbol: str):
    """Every call of one symbol, with speaker turns. Streams the file in batches."""
    import pyarrow as pa
    import pyarrow.compute as pc
    import pyarrow.parquet as pq

    cols = _META_COLS + ["structured_content"]
    parts = []
    for batch in pq.ParquetFile(path).iter_batches(batch_size=256, columns=cols):
        hit = pc.equal(batch.column("symbol"), symbol)
        if pc.any(hit).as_py():
            parts.append(batch.filter(hit))
    if not parts:
        return None
    df = pa.Table.from_batches(parts).to_pandas()
    df["date"] = df["date"].astype(str).str[:10]
    return df.sort_values(["date"], ascending=False).reset_index(drop=True)


# Older transcripts open with participant lists filed as "speakers".
_LIST_HEADS = {"executives", "analysts", "participants", "company participants",
               "corporate participants", "conference call participants", "management",
               "attendees", "presenters", "other participants"}
_OPERATOR_INTRO = re.compile(
    r"(?i)(my name is|i(?:'ll| will) be your|i am your|this is)\b.{0,80}\b"
    r"(operator|facilitator|moderator|coordinator)")


def _speaker(t: dict) -> str:
    return str(t.get("speaker", "")).strip().lower()


def qa_start(turns: list[dict]) -> int | None:
    """Index of the turn where the Q&A begins, or None if it cannot be found.

    A call opens with the operator (sometimes after participant lists), then the host
    and the prepared remarks; the first operator turn after a real speaker has
    spoken hands over to the questions ("[Operator Instructions]", "our first
    question comes from ..."). The operator is the speaker named Operator or, as in
    NVIDIA's calls, the first speaker if that turn introduces themselves as the
    conference operator. On a sample of 2,200 calls a boundary was found in 99% and
    the turn at it reads as a Q&A hand-over in 98.7% of those; the rest are
    operator interjections during the prepared remarks or empty turns, which is why
    the output names the turn it split at.
    """
    ops = {_speaker(t) for t in turns if _speaker(t).startswith("operator")}
    for t in turns[:4]:
        if _speaker(t) not in _LIST_HEADS and _OPERATOR_INTRO.search(str(t.get("text", ""))[:600]):
            ops.add(_speaker(t))
            break
    spoken = False
    for i, t in enumerate(turns):
        s = _speaker(t)
        if s in _LIST_HEADS or s.endswith("participants"):
            continue
        if s in ops:
            if spoken:
                return i
            continue
        spoken = True
    return None


def get(symbol: str, year: int | None = None, quarter: int | None = None) -> dict | None:
    """One call with its turns; the most recent when year/quarter are omitted."""
    p = _path()
    if not p:
        return None
    df = _calls(p, symbol.upper())
    if df is None:
        return None
    if year is not None:
        df = df[df["year"] == int(year)]
    if quarter is not None:
        df = df[df["quarter"] == int(quarter)]
    if df.empty:
        return None
    r = df.iloc[0]                       # newest by call date
    sc = r["structured_content"]
    turns = [{"speaker": t.get("speaker"), "text": t.get("text")}
             for t in (list(sc) if sc is not None else [])]
    return {"symbol": r["symbol"], "company": r["company_name"], "year": int(r["year"]),
            "quarter": int(r["quarter"]), "date": r["date"], "turns": turns,
            "matches": int(len(df)), "source": f"{REPO} (MIT)"}


def _boundary(turns: list[dict], start: int) -> str:
    """The turn the split was made at, quoted so a wrong split is visible."""
    t = turns[start]
    return (f"Q&A taken to start at turn {start + 1} ({t.get('speaker')}: "
            f"\"{re.sub(r'\s+', ' ', str(t.get('text', '')))[:80]}\")")


def render(call: dict, section: str = "full") -> tuple[str | None, str]:
    """(text, note) for one section of a call. text is None, with the reason in
    note, when the requested section cannot be separated."""
    turns = call["turns"]
    if not turns:
        return None, "the dataset holds no speaker-tagged text for this call"
    start = qa_start(turns)
    if section == "full":
        picked, note = turns, ""
    elif section not in ("prepared_remarks", "qa"):
        return None, f"unknown section '{section}'; one of full | prepared_remarks | qa"
    elif start is None:
        return None, ("the Q&A boundary could not be found in this transcript (no "
                      "operator turn after the opening); use section='full'")
    elif section == "prepared_remarks":
        picked = turns[:start]
        note = f"turns 1-{start} of {len(turns)}; {_boundary(turns, start)}"
    else:
        picked = turns[start:]
        note = f"turns {start + 1}-{len(turns)} of {len(turns)}; {_boundary(turns, start)}"
    text = "\n\n".join(f"{t['speaker']}: {re.sub(r'[ \t]+', ' ', str(t['text'])).strip()}"
                       for t in picked)
    return text, note