#!/usr/bin/env python3
"""Utility for priming LiveStore and flushing sample candles to parquet.

This script spins up the configured streaming data provider (local simulator by
default), keeps it running for a short window, and forces flushes so that
parquet partitions are materialized under ``data/parquet/live``.
"""

from __future__ import annotations

import argparse
import asyncio
import contextlib
import os
from pathlib import Path
from typing import Tuple

import sys

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from app.streaming.live_store import live_store
from app.streaming.data_provider import data_provider


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Prime LiveStore and flush parquet output for a symbol/timeframe")
    parser.add_argument("--symbol", default="XAUUSD", help="Symbol to stream (default: XAUUSD)")
    parser.add_argument("--timeframe", default="1m", help="Timeframe to stream (default: 1m)")
    parser.add_argument(
        "--duration",
        type=int,
        default=120,
        help="How long to keep the stream hot (seconds, default: 120)",
    )
    parser.add_argument(
        "--flush-interval",
        type=int,
        default=30,
        help="How often to flush to parquet during the run (seconds, default: 30)",
    )
    return parser.parse_args()


async def _maybe_consume(queue: asyncio.Queue) -> None:
    """Continuously drain the queue so it doesn't grow unbounded."""

    try:
        while True:
            item = await queue.get()
            if item is None:
                break
    except asyncio.CancelledError:
        pass


async def _prime_store(symbol: str, timeframe: str, duration: int, flush_interval: int) -> Tuple[int, Path]:
    symbol = symbol.upper().replace("/", "")
    timeframe = timeframe.lower()

    os.makedirs(live_store.root, exist_ok=True)
    target_dir = Path(live_store.root) / symbol / timeframe
    target_dir.mkdir(parents=True, exist_ok=True)

    queue, unsubscribe = await data_provider.subscribe(symbol, timeframe=timeframe)
    consumer_task = asyncio.create_task(_maybe_consume(queue))

    try:
        loop = asyncio.get_running_loop()
        start = loop.time()
        next_flush = flush_interval

        while True:
            elapsed = loop.time() - start
            if elapsed >= duration:
                break
            sleep_for = max(0.0, min(next_flush - elapsed, 1.0))
            await asyncio.sleep(sleep_for)
            elapsed = loop.time() - start
            if elapsed + 1e-6 >= next_flush:
                await live_store.flush_parquet()
                next_flush += flush_interval
    finally:
        await unsubscribe()
        consumer_task.cancel()
        with contextlib.suppress(Exception):
            await consumer_task

    await live_store.flush_parquet()
    history = live_store.get_history(symbol, timeframe)
    return len(history), target_dir


async def _async_main() -> None:
    args = _parse_args()
    bars, target_dir = await _prime_store(
        symbol=args.symbol,
        timeframe=args.timeframe,
        duration=args.duration,
        flush_interval=args.flush_interval,
    )

    partitions = sorted(target_dir.glob("date=*"))
    total_files = sum(len(list(part.glob("*.parquet"))) for part in partitions)
    latest_partition = partitions[-1] if partitions else None

    print("\nLiveStore setup complete:")
    print(f"  Symbol: {args.symbol.upper().replace('/', '')}")
    print(f"  Timeframe: {args.timeframe}")
    print(f"  Bars in memory: {bars}")
    print(f"  Root directory: {target_dir}")
    print(f"  Partitions found: {len(partitions)}")
    print(f"  Parquet files: {total_files}")
    if latest_partition:
        print(f"  Latest partition: {latest_partition}")


def main() -> None:
    asyncio.run(_async_main())


if __name__ == "__main__":
    main()