import datetime as dt
import sqlite3

import pandas as pd
import pytest

from app.screener.engine import ScreenContext
from app.screener.schema import ensure_screener_schema
from app.screener._test_fixtures import make_master, make_prices, make_flow


@pytest.fixture
def conn(tmp_path):
    db_path = tmp_path / "ctx.db"
    c = sqlite3.connect(db_path)
    ensure_screener_schema(c)
    yield c
    c.close()


def _seed(conn, master_df, prices_df, flow_df):
    now = dt.datetime.utcnow().isoformat()
    for t, row in master_df.iterrows():
        conn.execute("""INSERT INTO krx_master (ticker,name,market,market_cap,
            is_managed,is_preferred,is_spac,listed_date,updated_at)
            VALUES (?,?,?,?,?,?,?,?,?)""",
            (t, row["name"], row["market"], row["market_cap"],
             row["is_managed"], row["is_preferred"], row["is_spac"], None, now))
    prices_df.to_sql("krx_daily_prices", conn, if_exists="append", index=False)
    flow_df.to_sql("krx_flow", conn, if_exists="append", index=False)
    conn.commit()


def test_load_returns_dataframes(conn):
    asof = dt.date(2026, 5, 12)
    _seed(conn,
          make_master(["005930", "035420"]),
          make_prices(["005930", "035420"], days=30, asof=asof),
          make_flow(["005930", "035420"], days=30, asof=asof))

    ctx = ScreenContext.load(conn, asof, lookback_days=30)

    assert ctx.asof == asof
    assert set(ctx.master.index) == {"005930", "035420"}
    assert ctx.prices.shape[0] == 60                      # 2 종목 × 30일
    assert ctx.flow.shape[0] == 60


def test_restrict_filters_tickers(conn):
    asof = dt.date(2026, 5, 12)
    _seed(conn,
          make_master(["005930", "035420", "091990"]),
          make_prices(["005930", "035420", "091990"], days=30, asof=asof),
          make_flow(["005930", "035420", "091990"], days=30, asof=asof))

    ctx = ScreenContext.load(conn, asof, lookback_days=30)
    scoped = ctx.restrict(pd.Index(["005930"]))

    assert list(scoped.master.index) == ["005930"]
    assert (scoped.prices["ticker"] == "005930").all()
    assert (scoped.flow["ticker"] == "005930").all()