from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import pandas as pd
import sqlite3


@dataclass(frozen=True)
class KlineSource:
    db_path: Path
    table_name: str


def _to_ms(dt: datetime) -> int:
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    return int(dt.timestamp() * 1000)


def load_klines(
    source: KlineSource,
    start: datetime,
    end: datetime,
) -> pd.DataFrame:
    start_ms = _to_ms(start)
    end_ms = _to_ms(end)

    con = sqlite3.connect(str(source.db_path))
    try:
        df = pd.read_sql_query(
            f"SELECT id, open, high, low, close FROM {source.table_name} WHERE id >= ? AND id <= ? ORDER BY id ASC",
            con,
            params=(start_ms, end_ms),
        )
    finally:
        con.close()

    if df.empty:
        return df

    df["timestamp_ms"] = df["id"].astype("int64")
    df["dt"] = pd.to_datetime(df["timestamp_ms"], unit="ms", utc=True)
    df = df.drop(columns=["id"]).set_index("dt")

    for c in ("open", "high", "low", "close"):
        df[c] = pd.to_numeric(df[c], errors="coerce")

    df = df.dropna(subset=["open", "high", "low", "close"])
    return df