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, inspect
from sqlalchemy.orm import Session
from typing import Optional
from app.core.models._config import ServiceConfigValue


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()

    # Best-effort: ensure service_config table exists before waiting on it.
    # This complements the synchronous init in app.__main__ and protects
    # against ordering issues where _secrets is imported before that init runs.
    try:
        from sqlalchemy.exc import SQLAlchemyError

        with engine.begin() as conn:
            inspector = inspect(conn)
            if not inspector.has_table('service_config'):
                ServiceConfigValue.__table__.create(bind=conn, checkfirst=True)
    except SQLAlchemyError as exc:
        make_log("HotWallet", f"Failed to ensure service_config table: {exc}", level="error")
    except Exception:
        # Avoid failing hard here; the fallback waiter below may still succeed
        pass

    def db_ready(conn) -> bool:
        try:
            inspector = inspect(conn)
            return inspector.has_table('service_config')
        except Exception:
            return False

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

    def read_seed() -> Optional[bytes]:
        # Use a fresh connection/session per read to avoid snapshot staleness
        try:
            with engine.connect() as rconn:
                with Session(bind=rconn) as s:
                    row = s.query(ServiceConfigValue).filter(ServiceConfigValue.key == 'private_key').first()
                    if not row:
                        return None
                    packed = row.packed_value or {}
                    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:
            with engine.connect() as wconn:
                with Session(bind=wconn) as s:
                    s.add(ServiceConfigValue(key='private_key', packed_value={"value": seed.hex()}))
                    s.commit()
            make_log("HotWallet", "Seed saved in service_config by primary", level='info')
            return seed
        except Exception:
            # 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()