from os import getenv, urandom
import os
import time
import json

from nacl.bindings import crypto_sign_seed_keypair
from tonsdk.utils import Address

from app.core._blockchain.ton.wallet_v3cr3 import WalletV3CR3
from app.core.logger import make_log
from sqlalchemy import create_engine, text


def _load_seed_from_env_or_generate() -> bytes:
    seed_hex = os.getenv("TON_INIT_HOT_SEED")
    if seed_hex:
        make_log("HotWallet", "Loaded seed from env")
        return bytes.fromhex(seed_hex)
    make_log("HotWallet", "No seed provided; generating ephemeral seed", level='info')
    return urandom(32)


def _init_seed_via_db() -> bytes:
    """Store and read hot seed from PostgreSQL service_config (key='private_key').
    Primary node writes it once; workers wait until it appears.
    """
    from app.core._config import DATABASE_URL

    engine = create_engine(DATABASE_URL, pool_pre_ping=True)
    role = os.getenv("NODE_ROLE", "worker").lower()

    def db_ready(conn) -> bool:
        try:
            r = conn.execute(text("SELECT to_regclass('public.service_config')")).scalar()
            return r is not None
        except Exception:
            return False

    # Wait for table to exist
    start = time.time()
    with engine.connect() as conn:
        while not db_ready(conn):
            time.sleep(0.5)
            if time.time() - start > 120:
                raise TimeoutError("service_config table not available")

        def read_seed() -> bytes | None:
            row = conn.execute(text("SELECT packed_value FROM service_config WHERE key = :k LIMIT 1"), {"k": "private_key"}).first()
            if not row:
                return None
            packed = row[0] or {}
            try:
                # packed_value is JSON; ensure dict
                if isinstance(packed, str):
                    packed = json.loads(packed)
                seed_hex = packed.get("value")
                return bytes.fromhex(seed_hex) if seed_hex else None
            except Exception:
                return None

        seed = read_seed()
        if seed:
            return seed

        if role == "primary":
            seed = _load_seed_from_env_or_generate()
            # Try insert; if another primary raced, ignore
            try:
                conn.execute(
                    text("INSERT INTO service_config (key, packed_value) VALUES (:k, CAST(:v AS JSON))"),
                    {"k": "private_key", "v": json.dumps({"value": seed.hex()})}
                )
                conn.commit()
                make_log("HotWallet", "Seed saved in service_config by primary", level='info')
                return seed
            except Exception:
                conn.rollback()
                # Read again in case of race
                seed2 = read_seed()
                if seed2:
                    return seed2
                raise
        else:
            make_log("HotWallet", "Worker waiting for seed in service_config...", level='info')
            while True:
                seed = read_seed()
                if seed:
                    return seed
                time.sleep(0.5)


_extra_ton_wallet_options = {}
if getenv('TON_CUSTOM_WALLET_ADDRESS'):
    _extra_ton_wallet_options['address'] = Address(getenv('TON_CUSTOM_WALLET_ADDRESS'))


def _init_wallet():
    # Primary writes to DB; workers wait and read from DB
    hot_seed_bytes = _init_seed_via_db()
    pub, priv = crypto_sign_seed_keypair(hot_seed_bytes)
    wallet = WalletV3CR3(
        private_key=priv,
        public_key=pub,
        **_extra_ton_wallet_options
    )
    return hot_seed_bytes, pub, priv, wallet


hot_seed, hot_pubkey, hot_privkey, service_wallet = _init_wallet()