"""分批平仓:累计各批收回权利金后再算盈亏."""
from __future__ import annotations

import sqlite3
import unittest
from unittest.mock import patch

from lib.options.options_close_exec_lib import close_option_by_bid1
from lib.options.options_db import (
    add_open_close_accum,
    get_open_close_accum,
    init_options_tables,
)


def _mem_db() -> sqlite3.Connection:
    conn = sqlite3.connect(":memory:")
    conn.row_factory = sqlite3.Row
    init_options_tables(conn)
    return conn


class CloseAccumDbTests(unittest.TestCase):
    def test_accum_on_earliest_open(self):
        conn = _mem_db()
        conn.execute(
            """
            INSERT INTO options_trades
            (inst_id, underlying, opt_type, sheets, eth_amount, premium_paid, status)
            VALUES ('ETH-USD-260927-2700-P', 'ETH', 'P', 10, 0.1, 0.02, 'open')
            """
        )
        conn.commit()
        self.assertTrue(
            add_open_close_accum(
                conn, "ETH-USD-260927-2700-P", recv=0.005, sheets=5, px=0.05, ord_id="a1"
            )
        )
        self.assertTrue(
            add_open_close_accum(
                conn, "ETH-USD-260927-2700-P", recv=0.004, sheets=5, px=0.04, ord_id="a2"
            )
        )
        acc = get_open_close_accum(conn, "ETH-USD-260927-2700-P")
        self.assertAlmostEqual(float(acc["recv"]), 0.009, places=6)
        self.assertEqual(int(acc["sheets"]), 10)
        self.assertAlmostEqual(float(acc["notional"]), 0.05 * 5 + 0.04 * 5, places=6)
        self.assertEqual(acc["close_ord_id"], "a1,a2")


class TwoStepCloseExecTests(unittest.TestCase):
    def test_two_partial_closes_sum_premium_received(self):
        conn = _mem_db()
        inst = "ETH-USD-260927-2700-P"
        conn.execute(
            """
            INSERT INTO options_trades
            (inst_id, underlying, opt_type, strike, sheets, eth_amount,
             open_quote, premium_paid, status, premium_ccy, margin_mode)
            VALUES (?, 'ETH', 'P', 2700, 10, 0.1, 0.2, 0.02, 'open', 'ETH', 'coin')
            """,
            (inst,),
        )
        conn.commit()

        # 持仓:第一次 10→5,第二次 5→0
        pos_seq = [
            [{"instId": inst, "availPos": "10", "pos": "10", "mgnMode": "isolated", "idxPx": "2500"}],
            [{"instId": inst, "availPos": "5", "pos": "5", "mgnMode": "isolated", "idxPx": "2500"}],
            [{"instId": inst, "availPos": "5", "pos": "5", "mgnMode": "isolated", "idxPx": "2500"}],
            [],  # 第二次下单后已平
        ]
        pos_i = {"n": 0}

        def fetch_positions(_ex):
            i = pos_i["n"]
            pos_i["n"] = min(i + 1, len(pos_seq) - 1)
            return pos_seq[i]

        order_px = {"n": 0}

        def place_limit(_ex, **kwargs):
            order_px["n"] += 1
            px = 0.05 if order_px["n"] == 1 else 0.04
            return {"ok": True, "px": px, "data": {"ordId": f"ord{order_px['n']}"}}

        def quote(_ex, _inst):
            return {"ok": True, "tick_sz": "0.0001", "ct_mult": 0.01, "bid": 0.05, "index_px": 2500}

        def book(_ex, _inst, _n):
            # 每次只够吃一半深度,强制分两批
            return {"bids": [{"px": 0.05, "sz": 5}]} if order_px["n"] < 1 else {"bids": [{"px": 0.04, "sz": 5}]}

        cfg = {
            "get_db": lambda: _reuse_conn(conn),
            "quote_option_contract": quote,
            "fetch_option_positions": fetch_positions,
            "fetch_option_book_depth": book,
            "place_option_limit_order": place_limit,
            "td_mode": "isolated",
        }

        with patch("lib.options.options_close_exec_lib.time.sleep", return_value=None), patch(
            "lib.exchange.okx_options_lib.invalidate_option_positions_cache"
        ), patch(
            "lib.options.options_pricing_lib.fetch_option_mark_px", return_value=0.05
        ), patch(
            "lib.options.options_close_exec_lib._pos_close_refs", return_value=(0.05, 0.0)
        ):
            r1 = close_option_by_bid1(cfg, ex=object(), inst_id=inst, sheets=10)
            self.assertTrue(r1.get("ok"), r1)
            self.assertFalse(r1.get("fully_closed"), r1)
            open_row = conn.execute(
                "SELECT close_recv_accum, status, realized_pnl FROM options_trades WHERE inst_id=?",
                (inst,),
            ).fetchone()
            self.assertEqual(open_row["status"], "open")
            self.assertGreater(float(open_row["close_recv_accum"] or 0), 0)
            self.assertIsNone(open_row["realized_pnl"])

            # 第二次买一深度
            def book2(_ex, _inst, _n):
                return {"bids": [{"px": 0.04, "sz": 5}]}

            cfg["fetch_option_book_depth"] = book2
            r2 = close_option_by_bid1(cfg, ex=object(), inst_id=inst, sheets=5)
            self.assertTrue(r2.get("ok"), r2)
            self.assertTrue(r2.get("fully_closed"), r2)

        closed = conn.execute(
            """
            SELECT premium_paid, premium_received, realized_pnl, status, close_recv_accum
            FROM options_trades WHERE inst_id=?
            """,
            (inst,),
        ).fetchone()
        self.assertEqual(closed["status"], "closed")
        # ct_mult=0.01: 5张×0.05 + 5张×0.04 = 0.0025 + 0.002 = 0.0045
        self.assertAlmostEqual(float(closed["premium_received"]), 0.0045, places=6)
        self.assertAlmostEqual(float(closed["realized_pnl"]), 0.0045 - 0.02, places=6)
        self.assertEqual(float(closed["close_recv_accum"] or 0), 0.0)
        # 若只记第二批,recv 会是 0.002,pnl=-0.018;合计应为 -0.0155
        self.assertNotAlmostEqual(float(closed["premium_received"]), 0.002, places=6)


def _reuse_conn(conn: sqlite3.Connection) -> sqlite3.Connection:
    """get_db 每次返回同一连接;close 时 no-op,避免 :memory: 被关掉."""

    class _Proxy:
        def __getattr__(self, name):
            return getattr(conn, name)

        def close(self):
            return None

    return _Proxy()  # type: ignore[return-value]


if __name__ == "__main__":
    unittest.main()