from __future__ import annotations

import asyncio
import os
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Dict, List, Tuple, Any
import glob
import shutil

import pyarrow as pa
import pyarrow.parquet as pq


@dataclass
class _Series:
    bars: List[Dict[str, Any]]
    last_flushed_ts: int


class LiveStore:
    def __init__(self, max_bars: int = 5000, root: str = "data/parquet/live") -> None:
        self._series: Dict[Tuple[str, str], _Series] = {}
        self._max_bars = max_bars
        self._root = root
        self._lock = asyncio.Lock()

    @property
    def root(self) -> str:
        return self._root

    def _get_series(self, symbol: str, timeframe: str) -> _Series:
        key = (symbol, timeframe)
        s = self._series.get(key)
        if not s:
            s = _Series(bars=[], last_flushed_ts=0)
            self._series[key] = s
        return s

    def get_history(self, symbol: str, timeframe: str) -> List[Dict[str, Any]]:
        s = self._get_series(symbol, timeframe)
        return list(s.bars)

    def ingest_bar(self, symbol: str, timeframe: str, bar: Dict[str, Any]) -> None:
        s = self._get_series(symbol, timeframe)
        if s.bars and s.bars[-1]["time"] == bar["time"]:
            # update last
            last = s.bars[-1]
            last["high"] = max(last["high"], bar["high"])
            last["low"] = min(last["low"], bar["low"])
            last["close"] = bar["close"]
            last["volume"] = last.get("volume", 0.0) + bar.get("volume", 0.0)
        else:
            s.bars.append(bar)
            if len(s.bars) > self._max_bars:
                s.bars.pop(0)

    async def flush_parquet(self) -> None:
        # Write new bars since last flush, partitioned by date
        async with self._lock:
            for (symbol, timeframe), s in self._series.items():
                new_rows = [b for b in s.bars if b["time"] > s.last_flushed_ts]
                if not new_rows:
                    continue
                # Partition by date
                rows_by_date: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
                for r in new_rows:
                    dt = datetime.utcfromtimestamp(int(r["time"]))
                    rows_by_date[dt.strftime("%Y-%m-%d")].append(r)
                for date_str, rows in rows_by_date.items():
                    table = pa.Table.from_pylist(rows)
                    base = os.path.join(self._root, symbol, timeframe)
                    out = os.path.join(base, f"date={date_str}")
                    os.makedirs(out, exist_ok=True)
                    # write one file per flush to this partition
                    pq.write_table(table, os.path.join(out, f"part-{int(datetime.utcnow().timestamp())}.parquet"))
                s.last_flushed_ts = max(b["time"] for b in new_rows)


# Singleton store
live_store = LiveStore()


async def periodic_flush(interval_sec: int = 60):
    while True:
        try:
            await live_store.flush_parquet()
        except Exception:
            pass
        await asyncio.sleep(interval_sec)


def _iter_partitions(root: str):
    """Yield (symbol, timeframe, partition_path, date_str) for existing partitions."""
    # root/symbol/timeframe/date=YYYY-MM-DD
    for sym_dir in glob.glob(f"{root}/*"):
        if not os.path.isdir(sym_dir):
            continue
        symbol = os.path.basename(sym_dir)
        for tf_dir in glob.glob(f"{sym_dir}/*"):
            if not os.path.isdir(tf_dir):
                continue
            timeframe = os.path.basename(tf_dir)
            for part_dir in glob.glob(f"{tf_dir}/date=*" ):
                if not os.path.isdir(part_dir):
                    continue
                date_str = os.path.basename(part_dir).split("=", 1)[-1]
                yield (symbol, timeframe, part_dir, date_str)


def prune_old_partitions(root: str, retention_days: int = 7) -> int:
    """Delete partition directories older than retention_days. Returns count deleted."""
    now = datetime.utcnow()
    deleted = 0
    for symbol, timeframe, part_dir, date_str in list(_iter_partitions(root)):
        try:
            y, m, d = map(int, date_str.split("-"))
            dt = datetime(y, m, d)
            if now - dt > timedelta(days=retention_days):
                shutil.rmtree(part_dir, ignore_errors=True)
                deleted += 1
        except Exception:
            # Skip unparsable date partitions
            continue
    return deleted


def compact_partition(part_dir: str, max_files_threshold: int = 20) -> bool:
    """If too many small part files exist, compact them into a single file.
    Returns True if compaction performed.
    """
    part_files = sorted(glob.glob(os.path.join(part_dir, "part-*.parquet")))
    if len(part_files) < max_files_threshold:
        return False
    try:
        tables: List[pa.Table] = []
        for p in part_files:
            tables.append(pq.read_table(p))
        if not tables:
            return False
        combined = pa.concat_tables(tables, promote=True)
        out_file = os.path.join(part_dir, f"compact-{int(datetime.utcnow().timestamp())}.parquet")
        pq.write_table(combined, out_file)
        # remove old parts
        for p in part_files:
            try:
                os.remove(p)
            except Exception:
                pass
        return True
    except Exception:
        return False


def compact_all(root: str, max_files_threshold: int = 20) -> int:
    """Run compaction across all partitions. Returns number of partitions compacted."""
    compacted = 0
    for _, _, part_dir, _ in list(_iter_partitions(root)):
        if compact_partition(part_dir, max_files_threshold=max_files_threshold):
            compacted += 1
    return compacted


async def periodic_maintenance(retention_days: int = 7, compact_threshold_files: int = 20, interval_sec: int = 900):
    """Periodically prune old partitions and compact small files."""
    while True:
        try:
            prune_old_partitions(live_store.root, retention_days=retention_days)
            compact_all(live_store.root, max_files_threshold=compact_threshold_files)
        except Exception:
            pass
        await asyncio.sleep(interval_sec)