storage/briefs_db.py

"""storage/briefs_db.py — persist the last generated brief per ticker.

Stores the full BriefOutput dict as JSON, keyed by ticker.
One row per ticker: re-generating overwrites the previous brief.
Powers the sidebar Coverage list and brief reload across sessions
without re-running the LLM agent.
"""
from __future__ import annotations

import json
import hashlib
import sqlite3
from pathlib import Path
from typing import Optional

DB_PATH = Path("data/briefs.db")
BRIEF_SCHEMA_VERSION = "brief.v2"


def init_db() -> None:
    """Create the briefs table if it doesn't exist. Safe to call multiple times."""
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    with sqlite3.connect(DB_PATH) as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS briefs (
                ticker      TEXT PRIMARY KEY,
                filing_date TEXT,
                brief_json  TEXT NOT NULL,
                saved_at    TEXT NOT NULL
            )
        """)

        # Migrate language column if absent
        columns = conn.execute("PRAGMA table_info(briefs)").fetchall()
        column_names = [col[1] for col in columns]
        if "language" not in column_names:
            conn.execute(
                "ALTER TABLE briefs ADD COLUMN language TEXT NOT NULL DEFAULT 'English'"
            )

        # Create translations table
        conn.execute("""
            CREATE TABLE IF NOT EXISTS brief_translations (
                ticker TEXT NOT NULL,
                language TEXT NOT NULL,
                brief_json TEXT NOT NULL,
                saved_at TEXT NOT NULL,
                PRIMARY KEY (ticker, language)
            )
        """)

        # Append-only snapshots are the authoritative cache.  The legacy
        # ticker-keyed table remains as a compatibility index for old installs.
        conn.execute("""
            CREATE TABLE IF NOT EXISTS brief_snapshots (
                snapshot_id   TEXT PRIMARY KEY,
                ticker        TEXT NOT NULL,
                filing_date   TEXT,
                data_as_of    TEXT,
                generated_at  TEXT NOT NULL,
                schema_version TEXT NOT NULL,
                language      TEXT NOT NULL,
                status        TEXT NOT NULL,
                brief_json    TEXT NOT NULL
            )
        """)
        conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_brief_snapshots_latest
            ON brief_snapshots (ticker, generated_at DESC)
        """)


def save_brief(ticker: str, brief_dict: dict) -> None:
    """Persist an immutable snapshot and refresh the legacy latest-ticker index."""
    from datetime import datetime, timezone
    init_db()
    ticker = ticker.upper()
    language = brief_dict.get("language") or "English"
    generated_at = brief_dict.get("generated_at") or datetime.now(timezone.utc).isoformat()
    filing_date = brief_dict.get("filing_date", "")
    data_as_of = brief_dict.get("data_as_of") or filing_date
    schema_version = brief_dict.get("schema_version") or BRIEF_SCHEMA_VERSION
    status = brief_dict.get("status") or "PARTIAL"
    digest = hashlib.sha256(
        "\x1f".join((ticker, filing_date, generated_at, schema_version, language)).encode("utf-8")
    ).hexdigest()[:32]
    payload = json.dumps(brief_dict)
    with sqlite3.connect(DB_PATH) as conn:
        conn.execute(
            """
            INSERT OR IGNORE INTO brief_snapshots (
                snapshot_id, ticker, filing_date, data_as_of, generated_at,
                schema_version, language, status, brief_json
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                digest, ticker, filing_date, data_as_of, generated_at,
                schema_version, language, status, payload,
            ),
        )
        conn.execute(
            """
            INSERT INTO briefs (ticker, filing_date, brief_json, language, saved_at)
            VALUES (?, ?, ?, ?, ?)
            ON CONFLICT(ticker) DO UPDATE SET
                filing_date = excluded.filing_date,
                brief_json  = excluded.brief_json,
                language    = excluded.language,
                saved_at    = excluded.saved_at
            """,
            (
                ticker,
                filing_date,
                payload,
                language,
                generated_at,
            ),
        )
        # Invalidate stale translations for this ticker
        conn.execute("DELETE FROM brief_translations WHERE ticker = ?", (ticker,))


def get_brief(
    ticker: str,
    filing_date: Optional[str] = None,
    as_of: Optional[str] = None,
    snapshot_id: Optional[str] = None,
    *,
    include_metadata: bool = False,
) -> Optional[dict]:
    """Return a snapshot, latest by default, with optional point-in-time filters."""
    if not DB_PATH.exists():
        return None
    init_db()
    with sqlite3.connect(DB_PATH) as conn:
        clauses = ["ticker = ?"]
        params: list[str] = [ticker.upper()]
        if filing_date is not None:
            clauses.append("filing_date = ?")
            params.append(filing_date)
        if as_of is not None:
            clauses.append("generated_at <= ?")
            params.append(as_of)
        if snapshot_id is not None:
            clauses.append("snapshot_id = ?")
            params.append(snapshot_id)
        row = conn.execute(
            "SELECT brief_json, language, snapshot_id, filing_date, data_as_of, "
            "generated_at, schema_version, status FROM brief_snapshots WHERE "
            + " AND ".join(clauses)
            + " ORDER BY generated_at DESC LIMIT 1",
            params,
        ).fetchone()
        legacy = False
        if row is None and filing_date is None and as_of is None and snapshot_id is None:
            row = conn.execute(
                "SELECT brief_json, language, NULL, filing_date, filing_date, "
                "saved_at, 'legacy.v1', 'PARTIAL' FROM briefs WHERE ticker = ?",
                (ticker.upper(),),
            ).fetchone()
            legacy = row is not None
    if row is None:
        return None
    try:
        brief = json.loads(row[0])
        brief.setdefault("language", row[1] or "English")
        if legacy:
            brief["data_quality_status"] = "LEGACY_UNVERIFIED"
            brief["status"] = "PARTIAL"
            brief["display_policy"] = {
                "event_returns_aligned": False,
                "market_expectations_aligned": False,
                "aggregate_reliability_meaningful": False,
            }
            market = brief.get("market_expectations")
            if isinstance(market, dict):
                for field in (
                    "consensus_eps_est", "consensus_rev_est_bn", "revision_30d_pct",
                    "d1_price_reaction_pct", "d5_price_reaction_pct",
                    "since_release_price_reaction_pct",
                ):
                    market[field] = None
                market["period_aligned"] = False
                market["comparison_allowed"] = False
                market["event_aligned"] = False
                market["event_comparison_allowed"] = False
        if include_metadata:
            brief["_snapshot"] = {
                "snapshot_id": row[2],
                "filing_date": row[3] or "",
                "data_as_of": row[4] or "",
                "generated_at": row[5] or "",
                "schema_version": row[6],
                "status": row[7],
                "legacy": legacy,
            }
        return brief
    except json.JSONDecodeError:
        import sys
        print(f"[briefs_db] corrupt JSON for ticker {ticker.upper()!r}", file=sys.stderr)
        return None


def list_tickers() -> list[str]:
    """Return all tickers that have a persisted brief, sorted alphabetically."""
    if not DB_PATH.exists():
        return []
    init_db()
    with sqlite3.connect(DB_PATH) as conn:
        rows = conn.execute(
            "SELECT ticker FROM briefs UNION SELECT ticker FROM brief_snapshots ORDER BY ticker ASC"
        ).fetchall()
    return [r[0] for r in rows]


def delete_brief(ticker: str) -> int:
    """Delete a ticker's legacy row, snapshots, and cached translations."""
    init_db()
    ticker = ticker.upper()
    deleted = 0
    with sqlite3.connect(DB_PATH) as conn:
        for table in ("briefs", "brief_snapshots", "brief_translations"):
            cursor = conn.execute(f"DELETE FROM {table} WHERE ticker = ?", (ticker,))
            deleted += cursor.rowcount
    return deleted


def list_briefs() -> list[dict]:
    """Return one row per persisted brief: {ticker, filing_date, saved_at}.

    Sorted alphabetically by ticker. Powers the sidebar Coverage list.
    """
    if not DB_PATH.exists():
        return []
    init_db()
    with sqlite3.connect(DB_PATH) as conn:
        rows = conn.execute("""
            SELECT ticker, filing_date, generated_at
            FROM brief_snapshots AS s
            WHERE generated_at = (
                SELECT MAX(generated_at) FROM brief_snapshots
                WHERE ticker = s.ticker
            )
            ORDER BY ticker ASC
        """).fetchall()
        snapshot_tickers = {r[0] for r in rows}
        legacy_rows = conn.execute(
            "SELECT ticker, filing_date, saved_at FROM briefs ORDER BY ticker ASC"
        ).fetchall()
        rows.extend(r for r in legacy_rows if r[0] not in snapshot_tickers)
        rows.sort(key=lambda r: r[0])
    return [
        {"ticker": r[0], "filing_date": r[1] or "", "saved_at": r[2] or ""}
        for r in rows
    ]


def list_snapshots(ticker: str) -> list[dict]:
    """Return immutable snapshot metadata, newest first, without loading payloads."""
    if not DB_PATH.exists():
        return []
    init_db()
    with sqlite3.connect(DB_PATH) as conn:
        rows = conn.execute(
            """
            SELECT snapshot_id, filing_date, data_as_of, generated_at,
                   schema_version, language, status
            FROM brief_snapshots
            WHERE ticker = ?
            ORDER BY generated_at DESC
            """,
            (ticker.upper(),),
        ).fetchall()
    keys = (
        "snapshot_id", "filing_date", "data_as_of", "generated_at",
        "schema_version", "language", "status",
    )
    return [dict(zip(keys, row)) for row in rows]


def get_previous_brief(ticker: str, current_generated_at: Optional[str] = None) -> Optional[dict]:
    """Return the snapshot immediately before the current one, or None.

    ``brief_snapshots`` is append-only (see save_brief), so this is a pure
    lookup — no separate "history" table was needed. If
    *current_generated_at* is omitted, the latest snapshot's timestamp is
    used, so this returns the second-most-recent snapshot.
    """
    snapshots = list_snapshots(ticker)
    if len(snapshots) < 2:
        return None
    if current_generated_at is None:
        current_generated_at = snapshots[0]["generated_at"]
    for snap in snapshots:
        if snap["generated_at"] < current_generated_at:
            return get_brief(ticker, snapshot_id=snap["snapshot_id"])
    return None


def get_translation(ticker: str, language: str) -> Optional[dict]:
    """Return the persisted translation dict for ticker/language, or None if not found."""
    if not DB_PATH.exists():
        return None
    init_db()
    with sqlite3.connect(DB_PATH) as conn:
        row = conn.execute(
            "SELECT brief_json FROM brief_translations WHERE ticker = ? AND language = ?",
            (ticker.upper(), language),
        ).fetchone()
    if row is None:
        return None
    try:
        return json.loads(row[0])
    except json.JSONDecodeError:
        import sys
        print(
            f"[briefs_db] corrupt JSON for translation {ticker.upper()!r} ({language})",
            file=sys.stderr,
        )
        return None


def save_translation(ticker: str, language: str, brief_dict: dict) -> None:
    """Persist translated brief_dict for ticker/language, overwriting any previous translation."""
    from datetime import datetime, timezone
    init_db()
    with sqlite3.connect(DB_PATH) as conn:
        conn.execute(
            """
            INSERT OR REPLACE INTO brief_translations
            (ticker, language, brief_json, saved_at)
            VALUES (?, ?, ?, ?)
            """,
            (
                ticker.upper(),
                language,
                json.dumps(brief_dict),
                datetime.now(timezone.utc).isoformat(),
            ),
        )