#!/usr/bin/env python3
import hashlib
import json
import os
import re
import sqlite3
import sys
import time
import urllib.error
import urllib.request

import psycopg
from psycopg.types.json import Jsonb

DB_PATH = os.getenv("USAGE_DB_PATH", "/data/usage.sqlite")
PG_DSN = os.getenv("USAGE_SYNC_PG_DSN") or os.getenv("PGSTORE_DSN", "")
SCHEMA = os.getenv("USAGE_SYNC_SCHEMA", "cpa_manager_usage")
INTERVAL = int(os.getenv("USAGE_SYNC_INTERVAL_SEC", "3"))
BATCH = int(os.getenv("USAGE_SYNC_BATCH", "500"))
MANAGER_BASE = os.getenv("USAGE_SYNC_MANAGER_BASE", "http://127.0.0.1:18317").rstrip("/")
WAIT_SEC = int(os.getenv("USAGE_SYNC_WAIT_SEC", "300"))
FORCE_UPSTREAM_URL = os.getenv("USAGE_SYNC_FORCE_UPSTREAM_URL", "").strip()
FORCE_MANAGEMENT_KEY = os.getenv("USAGE_SYNC_FORCE_MANAGEMENT_KEY", "").strip()
FORCE_QUEUE = os.getenv("USAGE_SYNC_FORCE_QUEUE", "").strip()
FORCE_POP_SIDE = os.getenv("USAGE_SYNC_FORCE_POP_SIDE", "").strip()

if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", SCHEMA):
    raise SystemExit(f"invalid USAGE_SYNC_SCHEMA: {SCHEMA}")

QSCHEMA = SCHEMA


def log(msg: str):
    print(f"[usage-sync] {msg}", flush=True)


def now_ms() -> int:
    return int(time.time() * 1000)


def json_dumps(obj) -> str:
    return json.dumps(obj, ensure_ascii=False, sort_keys=True, separators=(",", ":"))


def row_hash(obj) -> str:
    return hashlib.sha256(json_dumps(obj).encode("utf-8")).hexdigest()


def pg_connect():
    if not PG_DSN:
        return None
    return psycopg.connect(PG_DSN, autocommit=True, prepare_threshold=None)


def sqlite_connect():
    conn = sqlite3.connect(DB_PATH, timeout=30, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA busy_timeout = 5000")
    conn.execute("PRAGMA journal_mode = WAL")
    return conn


def table_exists_sqlite(conn, table: str) -> bool:
    row = conn.execute(
        "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?",
        (table,),
    ).fetchone()
    return row is not None


def get_sqlite_columns(conn, table: str):
    rows = conn.execute(f"PRAGMA table_info({table})").fetchall()
    return [r["name"] for r in rows]


def fetch_sqlite_rows(conn, sql: str, params=()):
    return [dict(r) for r in conn.execute(sql, params).fetchall()]


def sqlite_count(conn, table: str) -> int:
    row = conn.execute(f"SELECT COUNT(*) AS c FROM {table}").fetchone()
    return int(row["c"] if row else 0)


def pg_get_state(pg, key: str, default="0"):
    row = pg.execute(
        f"SELECT value FROM {QSCHEMA}.sync_state WHERE key=%s",
        (key,),
    ).fetchone()
    return row[0] if row else default


def pg_set_state(pg, key: str, value: str):
    pg.execute(
        f"""
        INSERT INTO {QSCHEMA}.sync_state (key, value, updated_at_ms)
        VALUES (%s, %s, %s)
        ON CONFLICT (key) DO UPDATE
        SET value = EXCLUDED.value,
            updated_at_ms = EXCLUDED.updated_at_ms
        """,
        (key, str(value), now_ms()),
    )


def ensure_pg_schema(pg):
    pg.execute(f"CREATE SCHEMA IF NOT EXISTS {QSCHEMA}")

    pg.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {QSCHEMA}.settings (
            key text PRIMARY KEY,
            value_text text,
            row_json jsonb NOT NULL,
            synced_at_ms bigint NOT NULL
        )
        """
    )

    pg.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {QSCHEMA}.model_prices (
            model_key text PRIMARY KEY,
            row_json jsonb NOT NULL,
            synced_at_ms bigint NOT NULL
        )
        """
    )

    pg.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {QSCHEMA}.usage_events (
            event_hash text PRIMARY KEY,
            sqlite_id bigint,
            timestamp_ms bigint,
            model text,
            auth_index text,
            failed boolean,
            row_json jsonb NOT NULL,
            synced_at_ms bigint NOT NULL
        )
        """
    )
    pg.execute(
        f"CREATE INDEX IF NOT EXISTS idx_{QSCHEMA}_usage_events_ts ON {QSCHEMA}.usage_events(timestamp_ms DESC)"
    )

    pg.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {QSCHEMA}.dead_letter_events (
            record_hash text PRIMARY KEY,
            sqlite_id bigint,
            created_at_ms bigint,
            row_json jsonb NOT NULL,
            synced_at_ms bigint NOT NULL
        )
        """
    )

    pg.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {QSCHEMA}.sync_state (
            key text PRIMARY KEY,
            value text NOT NULL,
            updated_at_ms bigint NOT NULL
        )
        """
    )


def wait_http_ok(url: str, timeout_sec: int) -> bool:
    deadline = time.time() + timeout_sec
    while time.time() < deadline:
        try:
            with urllib.request.urlopen(url, timeout=5) as resp:
                if 200 <= resp.status < 300:
                    return True
        except Exception:
            pass
        time.sleep(2)
    return False


def wait_sqlite_ready(timeout_sec: int):
    deadline = time.time() + timeout_sec
    while time.time() < deadline:
        try:
            if not os.path.exists(DB_PATH):
                time.sleep(2)
                continue
            conn = sqlite_connect()
            ok = all(
                table_exists_sqlite(conn, t)
                for t in ("settings", "model_prices", "usage_events", "dead_letter_events")
            )
            conn.close()
            if ok:
                return True
        except Exception:
            pass
        time.sleep(2)
    return False


def normalize_setup_payload(raw: str):
    try:
        data = json.loads(raw)
    except Exception:
        log("PG 中 settings.setup 不是合法 JSON，跳过 setup 恢复")
        return None

    base = (
        FORCE_UPSTREAM_URL
        or data.get("cpaBaseUrl")
        or data.get("CPAUpstreamURL")
        or data.get("upstream")
        or ""
    ).strip().rstrip("/")

    key = (
        FORCE_MANAGEMENT_KEY
        or data.get("managementKey")
        or data.get("ManagementKey")
        or ""
    ).strip()

    queue = (
        FORCE_QUEUE
        or data.get("queue")
        or data.get("Queue")
        or "usage"
    ).strip()

    pop_side = (
        FORCE_POP_SIDE
        or data.get("popSide")
        or data.get("PopSide")
        or "right"
    ).strip()

    if not base or not key:
        log("PG 中 settings.setup 缺少 cpaBaseUrl/managementKey，跳过 setup 恢复")
        return None

    return {
        "cpaBaseUrl": base,
        "managementKey": key,
        "queue": queue,
        "popSide": pop_side,
    }


def restore_setup_via_http(pg):
    row = pg.execute(
        f"SELECT value_text FROM {QSCHEMA}.settings WHERE key='setup'"
    ).fetchone()
    if not row or not row[0]:
        log("PG 中还没有 setup，等待你首次在面板保存后再自动镜像")
        return False

    payload = normalize_setup_payload(row[0])
    if not payload:
        return False

    deadline = time.time() + WAIT_SEC
    body = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        f"{MANAGER_BASE}/setup",
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    while time.time() < deadline:
        try:
            with urllib.request.urlopen(req, timeout=10) as resp:
                if 200 <= resp.status < 300:
                    log(f"已从 PG 恢复 cpa-manager setup，upstream={payload['cpaBaseUrl']}")
                    return True
        except urllib.error.HTTPError as e:
            try:
                detail = e.read().decode("utf-8", "ignore")
            except Exception:
                detail = str(e)
            log(f"/setup 返回 {e.code}，稍后重试：{detail}")
        except Exception as e:
            log(f"/setup 暂时失败，稍后重试：{e}")
        time.sleep(2)

    log("setup 自动恢复超时；如果面板仍未接管统计，请手动打开 management.html 保存一次")
    return False


def sqlite_insert(conn, table: str, row: dict, mode: str):
    cols = get_sqlite_columns(conn, table)
    use_cols = [c for c in cols if c in row]
    if not use_cols:
        return
    placeholders = ",".join("?" for _ in use_cols)
    col_sql = ",".join(use_cols)
    values = [row[c] for c in use_cols]
    conn.execute(
        f"INSERT OR {mode} INTO {table} ({col_sql}) VALUES ({placeholders})",
        values,
    )


def restore_settings_to_sqlite(pg, conn):
    rows = pg.execute(
        f"SELECT row_json FROM {QSCHEMA}.settings WHERE key <> 'setup'"
    ).fetchall()
    for (row_json,) in rows:
        sqlite_insert(conn, "settings", dict(row_json), "REPLACE")
    conn.commit()


def restore_table_if_empty(pg, conn, table: str, order_sql: str):
    if sqlite_count(conn, table) > 0:
        return

    total = 0
    while True:
        rows = pg.execute(
            f"SELECT row_json FROM {QSCHEMA}.{table} {order_sql} LIMIT %s OFFSET %s",
            (BATCH, total),
        ).fetchall()
        if not rows:
            break

        for (row_json,) in rows:
            mode = "REPLACE" if table in ("settings", "model_prices") else "IGNORE"
            sqlite_insert(conn, table, dict(row_json), mode)
        conn.commit()
        total += len(rows)
        log(f"已从 PG 恢复 {table}: {total}")

    if total:
        log(f"{table} 恢复完成，共 {total} 条")


def restore_from_pg_if_needed(pg):
    if not wait_sqlite_ready(WAIT_SEC):
        log("SQLite 表长期未就绪，跳过 PG -> SQLite 恢复")
        return

    conn = sqlite_connect()
    try:
        restore_settings_to_sqlite(pg, conn)
        restore_table_if_empty(pg, conn, "model_prices", "ORDER BY model_key ASC")
        restore_table_if_empty(pg, conn, "usage_events", "ORDER BY COALESCE(sqlite_id, 0) ASC, event_hash ASC")
        restore_table_if_empty(pg, conn, "dead_letter_events", "ORDER BY COALESCE(sqlite_id, 0) ASC, record_hash ASC")

        max_usage = conn.execute("SELECT COALESCE(MAX(id), 0) FROM usage_events").fetchone()[0]
        max_dead = conn.execute("SELECT COALESCE(MAX(id), 0) FROM dead_letter_events").fetchone()[0]
        pg_set_state(pg, "sqlite_usage_last_id", str(max_usage))
        pg_set_state(pg, "sqlite_dead_last_id", str(max_dead))
    finally:
        conn.close()


def sync_settings_up(pg, conn):
    rows = fetch_sqlite_rows(conn, "SELECT * FROM settings")
    for row in rows:
        key = str(row.get("key", ""))
        if not key:
            continue
        value_text = row.get("value")
        pg.execute(
            f"""
            INSERT INTO {QSCHEMA}.settings (key, value_text, row_json, synced_at_ms)
            VALUES (%s, %s, %s, %s)
            ON CONFLICT (key) DO UPDATE
            SET value_text = EXCLUDED.value_text,
                row_json = EXCLUDED.row_json,
                synced_at_ms = EXCLUDED.synced_at_ms
            """,
            (key, value_text, Jsonb(row), now_ms()),
        )


def sync_model_prices_up(pg, conn):
    rows = fetch_sqlite_rows(conn, "SELECT * FROM model_prices")
    for row in rows:
        model_key = str(row.get("model") or row_hash(row))
        pg.execute(
            f"""
            INSERT INTO {QSCHEMA}.model_prices (model_key, row_json, synced_at_ms)
            VALUES (%s, %s, %s)
            ON CONFLICT (model_key) DO UPDATE
            SET row_json = EXCLUDED.row_json,
                synced_at_ms = EXCLUDED.synced_at_ms
            """,
            (model_key, Jsonb(row), now_ms()),
        )


def sync_usage_events_up(pg, conn):
    last_id = int(pg_get_state(pg, "sqlite_usage_last_id", "0"))
    rows = fetch_sqlite_rows(
        conn,
        "SELECT * FROM usage_events WHERE id > ? ORDER BY id ASC LIMIT ?",
        (last_id, BATCH),
    )
    if not rows:
        return 0

    max_id = last_id
    for row in rows:
        event_hash = str(row.get("event_hash") or row_hash(row))
        sqlite_id = int(row.get("id") or 0)
        timestamp_ms = int(row.get("timestamp_ms") or 0)
        model = row.get("model")
        auth_index = row.get("auth_index")
        failed = bool(row.get("failed") or False)

        pg.execute(
            f"""
            INSERT INTO {QSCHEMA}.usage_events
            (event_hash, sqlite_id, timestamp_ms, model, auth_index, failed, row_json, synced_at_ms)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT (event_hash) DO UPDATE
            SET sqlite_id = EXCLUDED.sqlite_id,
                timestamp_ms = EXCLUDED.timestamp_ms,
                model = EXCLUDED.model,
                auth_index = EXCLUDED.auth_index,
                failed = EXCLUDED.failed,
                row_json = EXCLUDED.row_json,
                synced_at_ms = EXCLUDED.synced_at_ms
            """,
            (
                event_hash,
                sqlite_id,
                timestamp_ms,
                model,
                auth_index,
                failed,
                Jsonb(row),
                now_ms(),
            ),
        )
        if sqlite_id > max_id:
            max_id = sqlite_id

    pg_set_state(pg, "sqlite_usage_last_id", str(max_id))
    return len(rows)


def sync_dead_letters_up(pg, conn):
    last_id = int(pg_get_state(pg, "sqlite_dead_last_id", "0"))
    rows = fetch_sqlite_rows(
        conn,
        "SELECT * FROM dead_letter_events WHERE id > ? ORDER BY id ASC LIMIT ?",
        (last_id, BATCH),
    )
    if not rows:
        return 0

    max_id = last_id
    for row in rows:
        sqlite_id = int(row.get("id") or 0)
        record_hash = str(
            row.get("record_hash")
            or row_hash(
                {
                    "payload": row.get("payload"),
                    "error": row.get("error"),
                    "created_at_ms": row.get("created_at_ms"),
                    "id": sqlite_id,
                }
            )
        )
        created_at_ms = int(row.get("created_at_ms") or 0)

        pg.execute(
            f"""
            INSERT INTO {QSCHEMA}.dead_letter_events
            (record_hash, sqlite_id, created_at_ms, row_json, synced_at_ms)
            VALUES (%s, %s, %s, %s, %s)
            ON CONFLICT (record_hash) DO UPDATE
            SET sqlite_id = EXCLUDED.sqlite_id,
                created_at_ms = EXCLUDED.created_at_ms,
                row_json = EXCLUDED.row_json,
                synced_at_ms = EXCLUDED.synced_at_ms
            """,
            (
                record_hash,
                sqlite_id,
                created_at_ms,
                Jsonb(row),
                now_ms(),
            ),
        )
        if sqlite_id > max_id:
            max_id = sqlite_id

    pg_set_state(pg, "sqlite_dead_last_id", str(max_id))
    return len(rows)


def sync_loop(pg):
    if not wait_sqlite_ready(WAIT_SEC):
        log("SQLite 长期未就绪，退出同步进程")
        sys.exit(1)

    while True:
        try:
            conn = sqlite_connect()
            try:
                sync_settings_up(pg, conn)
                sync_model_prices_up(pg, conn)

                moved_usage = 0
                while True:
                    n = sync_usage_events_up(pg, conn)
                    moved_usage += n
                    if n < BATCH:
                        break

                moved_dead = 0
                while True:
                    n = sync_dead_letters_up(pg, conn)
                    moved_dead += n
                    if n < BATCH:
                        break

                if moved_usage or moved_dead:
                    log(f"本轮已同步 usage_events={moved_usage}, dead_letter_events={moved_dead}")
            finally:
                conn.close()
        except Exception as e:
            log(f"同步失败：{e}")

        time.sleep(INTERVAL)


def main():
    if not PG_DSN:
        log("未设置 PGSTORE_DSN / USAGE_SYNC_PG_DSN，仅使用本地 SQLite，不做 PG 镜像")
        while True:
            time.sleep(3600)

    pg = pg_connect()
    ensure_pg_schema(pg)
    log(f"PG 镜像表已就绪，schema={SCHEMA}")

    if not wait_http_ok(f"{MANAGER_BASE}/health", WAIT_SEC):
        log("cpa-manager /health 长期未就绪，退出")
        sys.exit(1)

    restore_setup_via_http(pg)
    restore_from_pg_if_needed(pg)
    sync_loop(pg)


if __name__ == "__main__":
    main()