from __future__ import annotations

from datetime import datetime
from typing import Any


class DragonTigerRepositoryMixin:
    def list_seat_aliases(self) -> dict[str, str]:
        with self.connect() as connection:
            rows = connection.execute("SELECT seat_name, alias FROM seat_aliases").fetchall()
        return {row["seat_name"]: row["alias"] for row in rows}

    def save_seat_alias(self, seat_name: str, alias: str) -> None:
        now = datetime.now().astimezone().isoformat(timespec="seconds")
        with self.connect() as connection:
            connection.execute(
                """
                INSERT INTO seat_aliases (seat_name, alias, updated_at)
                VALUES (?, ?, ?)
                ON CONFLICT(seat_name) DO UPDATE SET
                    alias = excluded.alias,
                    updated_at = excluded.updated_at
                """,
                (seat_name, alias, now),
            )

    def upsert_lhb_institutions(self, rows: list[dict[str, Any]]) -> int:
        grouped: dict[tuple[str, str], dict[str, float | int]] = {}
        for row in rows:
            trade_date = str(row.get("trade_date") or "")
            ts_code = str(row.get("ts_code") or "")
            seat_name = str(row.get("exalter") or row.get("seat_name") or "")
            if not trade_date or not ts_code or "机构专用" not in seat_name:
                continue
            group = grouped.setdefault(
                (trade_date, ts_code),
                {"net": 0.0, "buy": 0.0, "sell": 0.0, "seats": 0},
            )
            group["net"] = float(group["net"]) + float(row.get("net_buy") or row.get("net_amount") or 0)
            group["buy"] = float(group["buy"]) + float(row.get("buy") or row.get("buy_amount") or 0)
            group["sell"] = float(group["sell"]) + float(row.get("sell") or row.get("sell_amount") or 0)
            group["seats"] = int(group["seats"]) + 1
        values = [
            (trade_date, ts_code, item["net"], item["buy"], item["sell"], item["seats"])
            for (trade_date, ts_code), item in grouped.items()
        ]
        with self.connect() as connection:
            connection.executemany(
                """
                INSERT INTO lhb_institution_daily
                    (trade_date, ts_code, net_buy_amount, buy_amount, sell_amount, seat_count)
                VALUES (?, ?, ?, ?, ?, ?)
                ON CONFLICT(trade_date, ts_code) DO UPDATE SET
                    net_buy_amount=excluded.net_buy_amount,
                    buy_amount=excluded.buy_amount,
                    sell_amount=excluded.sell_amount,
                    seat_count=excluded.seat_count
                """,
                values,
            )
        return len(values)
