data/institutional_flow_cache.py

"""Durable SQLite cache for Taiwan institutional investor flow rows."""

from __future__ import annotations

import sqlite3
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable

from data.institutional_flow import FLOW_COLUMNS

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_CACHE_PATH = ROOT / "data" / "cache" / "institutional_flow.sqlite"

DAILY_TABLE = "institutional_daily"
STATUS_TABLE = "institutional_fetch_status"

DAILY_COLUMNS = [
    "date",
    "stock_no",
    *FLOW_COLUMNS,
    "institutional_available",
    "institutional_source",
    "institutional_as_of",
    "institutional_carry_forward",
]


def connect(db_path: str | Path) -> sqlite3.Connection:
    """Open a SQLite connection with dict-like row access."""
    Path(db_path).parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    return conn


def init_db(db: str | Path | sqlite3.Connection) -> sqlite3.Connection:
    """Create cache tables if needed and return an initialized connection."""
    conn = db if isinstance(db, sqlite3.Connection) else connect(db)
    conn.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {DAILY_TABLE} (
            date TEXT NOT NULL,
            stock_no TEXT NOT NULL,
            foreign_buy INTEGER NOT NULL DEFAULT 0,
            foreign_sell INTEGER NOT NULL DEFAULT 0,
            foreign_net INTEGER NOT NULL DEFAULT 0,
            trust_buy INTEGER NOT NULL DEFAULT 0,
            trust_sell INTEGER NOT NULL DEFAULT 0,
            trust_net INTEGER NOT NULL DEFAULT 0,
            dealer_buy INTEGER NOT NULL DEFAULT 0,
            dealer_sell INTEGER NOT NULL DEFAULT 0,
            dealer_net INTEGER NOT NULL DEFAULT 0,
            institutional_buy INTEGER NOT NULL DEFAULT 0,
            institutional_sell INTEGER NOT NULL DEFAULT 0,
            institutional_net INTEGER NOT NULL DEFAULT 0,
            institutional_available INTEGER NOT NULL DEFAULT 1,
            institutional_source TEXT NOT NULL DEFAULT '',
            institutional_as_of TEXT,
            institutional_carry_forward INTEGER NOT NULL DEFAULT 0,
            updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
            PRIMARY KEY (stock_no, date)
        )
        """
    )
    conn.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {STATUS_TABLE} (
            source TEXT NOT NULL,
            date TEXT NOT NULL,
            status TEXT NOT NULL,
            fetched_at TEXT NOT NULL,
            row_count INTEGER NOT NULL DEFAULT 0,
            error TEXT NOT NULL DEFAULT '',
            PRIMARY KEY (source, date)
        )
        """
    )
    conn.execute(
        f"CREATE INDEX IF NOT EXISTS idx_{DAILY_TABLE}_date ON {DAILY_TABLE}(date)"
    )
    conn.commit()
    return conn


def _iso_date(value: Any) -> str:
    if hasattr(value, "isoformat"):
        return value.isoformat()
    return str(value).strip()


def _int_value(value: Any, default: int = 0) -> int:
    if value is None:
        return default
    if isinstance(value, bool):
        return int(value)
    try:
        return int(value)
    except (TypeError, ValueError):
        return default


def _row_values(row: dict[str, Any]) -> dict[str, Any]:
    values: dict[str, Any] = {
        "date": _iso_date(row["date"]),
        "stock_no": str(row["stock_no"]).replace(".TW", "").replace(".TWO", "").strip(),
        "institutional_available": _int_value(row.get("institutional_available", 1), 1),
        "institutional_source": str(row.get("institutional_source") or ""),
        "institutional_as_of": (
            _iso_date(row["institutional_as_of"])
            if row.get("institutional_as_of") is not None
            else None
        ),
        "institutional_carry_forward": _int_value(row.get("institutional_carry_forward", 0)),
    }
    for column in FLOW_COLUMNS:
        values[column] = _int_value(row.get(column, 0))
    return values


def upsert_daily_rows(
    db: str | Path | sqlite3.Connection,
    rows: Iterable[dict[str, Any]],
) -> int:
    """Insert or replace institutional daily rows; return rows written."""
    conn = init_db(db)
    prepared = [_row_values(row) for row in rows]
    if not prepared:
        return 0

    placeholders = ", ".join(f":{column}" for column in DAILY_COLUMNS)
    updates = ", ".join(
        f"{column}=excluded.{column}"
        for column in DAILY_COLUMNS
        if column not in {"date", "stock_no"}
    )
    conn.executemany(
        f"""
        INSERT INTO {DAILY_TABLE} ({", ".join(DAILY_COLUMNS)})
        VALUES ({placeholders})
        ON CONFLICT(stock_no, date) DO UPDATE SET
            {updates},
            updated_at=CURRENT_TIMESTAMP
        """,
        prepared,
    )
    conn.commit()
    return len(prepared)


def fetch_daily_rows(
    db: str | Path | sqlite3.Connection,
    stock_no: str,
    start_date: Any,
    end_date: Any,
) -> list[dict[str, Any]]:
    """Fetch cached institutional rows for one stock/date range, ascending."""
    conn = init_db(db)
    code = str(stock_no).replace(".TW", "").replace(".TWO", "").strip()
    cursor = conn.execute(
        f"""
        SELECT {", ".join(DAILY_COLUMNS)}
        FROM {DAILY_TABLE}
        WHERE stock_no = ? AND date >= ? AND date <= ?
        ORDER BY date ASC
        """,
        (code, _iso_date(start_date), _iso_date(end_date)),
    )
    return [_decode_daily_row(row) for row in cursor.fetchall()]


def fetch_recent_rows_by_code(
    db: str | Path | sqlite3.Connection = DEFAULT_CACHE_PATH,
    *,
    codes: Iterable[str] | None = None,
    lookback_days: int = 120,
    end_date: Any | None = None,
) -> dict[str, list[dict[str, Any]]]:
    """
    Return recent institutional rows grouped by stock code.

    This helper intentionally uses trading rows rather than calendar-day math:
    it first finds the latest available cache dates, then returns rows whose
    date is inside that latest-N trading-date set.  The hotspot scanner can use
    it without issuing TWSE/TPEx requests at request time.
    """
    conn = init_db(db)
    params: list[Any] = []
    where = ""
    if end_date is not None:
        where = "WHERE date <= ?"
        params.append(_iso_date(end_date))
    date_rows = conn.execute(
        f"""
        SELECT DISTINCT date
        FROM {DAILY_TABLE}
        {where}
        ORDER BY date DESC
        LIMIT ?
        """,
        (*params, max(1, int(lookback_days))),
    ).fetchall()
    dates = sorted(row["date"] for row in date_rows)
    if not dates:
        return {}

    code_list = [
        str(code).replace(".TW", "").replace(".TWO", "").strip()
        for code in (codes or [])
        if str(code).strip()
    ]
    query_params: list[Any] = [dates[0], dates[-1]]
    code_clause = ""
    if code_list:
        placeholders = ", ".join("?" for _ in code_list)
        code_clause = f" AND stock_no IN ({placeholders})"
        query_params.extend(code_list)

    cursor = conn.execute(
        f"""
        SELECT {", ".join(DAILY_COLUMNS)}
        FROM {DAILY_TABLE}
        WHERE date >= ? AND date <= ? {code_clause}
        ORDER BY stock_no ASC, date ASC
        """,
        query_params,
    )
    grouped: dict[str, list[dict[str, Any]]] = {}
    valid_dates = set(dates)
    for row in cursor.fetchall():
        decoded = _decode_daily_row(row)
        if decoded["date"] not in valid_dates:
            continue
        grouped.setdefault(decoded["stock_no"], []).append(decoded)
    return grouped


def _decode_daily_row(row: sqlite3.Row) -> dict[str, Any]:
    data = dict(row)
    data["institutional_available"] = bool(data["institutional_available"])
    data["institutional_carry_forward"] = bool(data["institutional_carry_forward"])
    return data


def upsert_fetch_status(
    db: str | Path | sqlite3.Connection,
    source: str,
    date: Any,
    status: str,
    *,
    row_count: int = 0,
    error: str = "",
    fetched_at: str | None = None,
) -> None:
    """Record the fetch outcome for a source/date pair."""
    conn = init_db(db)
    conn.execute(
        f"""
        INSERT INTO {STATUS_TABLE}
            (source, date, status, fetched_at, row_count, error)
        VALUES (?, ?, ?, ?, ?, ?)
        ON CONFLICT(source, date) DO UPDATE SET
            status=excluded.status,
            fetched_at=excluded.fetched_at,
            row_count=excluded.row_count,
            error=excluded.error
        """,
        (
            source,
            _iso_date(date),
            status,
            fetched_at or datetime.now(timezone.utc).isoformat(),
            _int_value(row_count),
            error or "",
        ),
    )
    conn.commit()


def get_fetch_status(
    db: str | Path | sqlite3.Connection,
    source: str,
    date: Any,
) -> dict[str, Any] | None:
    """Return the fetch status for a source/date pair, if present."""
    conn = init_db(db)
    row = conn.execute(
        f"""
        SELECT source, date, status, fetched_at, row_count, error
        FROM {STATUS_TABLE}
        WHERE source = ? AND date = ?
        """,
        (source, _iso_date(date)),
    ).fetchone()
    return dict(row) if row is not None else None