"""Extract historical bars for key liquid symbols into SQLite and JSON cache."""

from __future__ import annotations

import argparse
import json
import sqlite3
import time
from pathlib import Path

from drift_quant.nse_data.loader import (
    load_selected_symbols,
    save_selected_cache,
)

CORE_UNIVERSE = (
    # Banking
    "HDFCBANK",
    "ICICIBANK",
    "SBIN",
    "AXISBANK",
    "KOTAKBANK",
    "INDUSINDBK",
    "BANKBARODA",
    # Financial Services
    "BAJFINANCE",
    "BAJAJFINSV",
    "HDFCLIFE",
    "SBILIFE",
    # Information Technology
    "TCS",
    "INFY",
    "WIPRO",
    "HCLTECH",
    "TECHM",
    "LTIM",
    # Energy, Oil & Gas
    "RELIANCE",
    "ONGC",
    "BPCL",
    "POWERGRID",
    "NTPC",
    "TATAPOWER",
    "COALINDIA",
    # FMCG & Consumer
    "ITC",
    "HINDUNILVR",
    "NESTLEIND",
    "BRITANNIA",
    "DABUR",
    "TATACONSUM",
    # Automobile
    "TATAMOTORS",
    "MARUTI",
    "M&M",
    "BAJAJ-AUTO",
    "EICHERMOT",
    "HEROMOTOCO",
    # Metals & Mining
    "TATASTEEL",
    "JSWSTEEL",
    "HINDALCO",
    "VEDL",
    "JINDALSTEL",
    # Pharma & Healthcare
    "SUNPHARMA",
    "DRREDDY",
    "CIPLA",
    "DIVISLAB",
    "APOLLOHOSP",
    # Infrastructure, Industrials & Conglomerates
    "LT",
    "BHARTIARTL",
    "ULTRACEMCO",
    "GRASIM",
    "ADANIENT",
    "ADANIPORTS",
    "TITAN",
    "ASIANPAINT",
)


def populate_sqlite(dataset, db_path: Path) -> None:
    db_path.parent.mkdir(parents=True, exist_ok=True)
    if db_path.is_file():
        db_path.unlink()
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()
    cur.execute(
        """
        CREATE TABLE bars (
            symbol TEXT NOT NULL,
            timestamp TEXT NOT NULL,
            open REAL NOT NULL,
            high REAL NOT NULL,
            low REAL NOT NULL,
            close REAL NOT NULL,
            volume REAL NOT NULL,
            source_row INTEGER NOT NULL
        )
        """
    )
    records = []
    for _symbol, bars in dataset.bars_by_symbol.items():
        for bar in bars:
            records.append(
                (
                    bar.symbol,
                    bar.timestamp.isoformat(sep=" "),
                    bar.open,
                    bar.high,
                    bar.low,
                    bar.close,
                    bar.volume,
                    bar.source_row,
                )
            )
    cur.executemany(
        """
        INSERT INTO bars (symbol, timestamp, open, high, low, close, volume, source_row)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        """,
        records,
    )
    cur.execute("CREATE INDEX idx_bars_symbol_ts ON bars (symbol, timestamp)")
    conn.commit()
    conn.close()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "archive",
        type=Path,
        nargs="?",
        default=Path("all_NSE500_15min_data_2019_onwards.zip"),
    )
    parser.add_argument("--audit", type=Path, default=Path("artifacts/dataset_audit.json"))
    parser.add_argument(
        "--json-cache", type=Path, default=Path("artifacts/cache/nse_universe_20.json")
    )
    parser.add_argument(
        "--sqlite-cache", type=Path, default=Path("artifacts/cache/market_data.sqlite")
    )
    args = parser.parse_args()

    audit = json.loads(args.audit.read_text(encoding="utf-8"))
    sha = audit["archive_sha256"]

    print(f"Scanning {args.archive} for {len(CORE_UNIVERSE)} symbols across 2019-2025...")
    t0 = time.time()
    dataset = load_selected_symbols(
        args.archive,
        sha,
        CORE_UNIVERSE,
        progress=print,
    )
    elapsed = time.time() - t0
    print(
        f"Extracted {dataset.rows_selected:,} bars across {len(dataset.symbols)} symbols "
        f"in {elapsed:.1f}s."
    )

    print(f"Writing JSON cache to {args.json_cache}...")
    save_selected_cache(dataset, args.json_cache)

    print(f"Writing SQLite database to {args.sqlite_cache}...")
    populate_sqlite(dataset, args.sqlite_cache)

    print("Market data caching complete.")


if __name__ == "__main__":
    main()