"""每日 simulation 配額（5000 / 帳號）— EST-anchored，atomic SQLite upsert。

抄自 zhutoutoutousan/worldquant-miner generation_two/core/simulation_counter.py。
"""

from __future__ import annotations

import os
import sqlite3
from datetime import datetime
from typing import Tuple
from zoneinfo import ZoneInfo

_NY = ZoneInfo("America/New_York")

PATH = os.path.expanduser("~/.wq_gui_sim_counter.db")
DAILY_HARD = 5000
DAILY_WARN = 4000


def est_today() -> str:
    """以 America/New_York 時區回傳 'YYYY-MM-DD'，正確處理 EST/EDT 切換。"""
    return datetime.now(_NY).strftime("%Y-%m-%d")


def _init_db() -> None:
    with sqlite3.connect(PATH, timeout=5.0) as conn:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS sim_count (
                date_est TEXT PRIMARY KEY,
                count INTEGER NOT NULL DEFAULT 0
            )
        """
        )


def get_count() -> int:
    try:
        _init_db()
        with sqlite3.connect(PATH, timeout=5.0) as conn:
            cur = conn.execute(
                "SELECT count FROM sim_count WHERE date_est=?", (est_today(),)
            )
            row = cur.fetchone()
        return int(row[0]) if row else 0
    except Exception:
        return 0


def increment(n: int = 1) -> int:
    """atomic upsert。worker thread 安全，回傳累加後值。"""
    try:
        _init_db()
        today = est_today()
        with sqlite3.connect(PATH, timeout=5.0) as conn:
            conn.execute(
                """
                INSERT INTO sim_count(date_est, count) VALUES(?, ?)
                ON CONFLICT(date_est) DO UPDATE SET count = count + excluded.count
            """,
                (today, n),
            )
            cur = conn.execute("SELECT count FROM sim_count WHERE date_est=?", (today,))
            row = cur.fetchone()
        return int(row[0]) if row else n
    except Exception:
        return 0


def can_proceed() -> Tuple[bool, str]:
    cur = get_count()
    if cur >= DAILY_HARD:
        return False, f"今日 simulation 已達 {cur}/{DAILY_HARD} 上限"
    if cur >= DAILY_WARN:
        return True, f"⚠️ 今日 simulation {cur}/{DAILY_HARD} 接近上限"
    return True, f"今日 sim {cur}/{DAILY_HARD}"
