#!/usr/bin/env python3
import io
import json
import logging
import os
import re
import subprocess
import tempfile
from decimal import Decimal, InvalidOperation

from telegram import InlineKeyboardButton, InlineKeyboardMarkup, ReplyKeyboardMarkup, Update
from telegram.constants import ParseMode
from telegram.ext import (
    Application,
    CallbackQueryHandler,
    CommandHandler,
    ConversationHandler,
    ContextTypes,
    MessageHandler,
    filters,
)

logging.basicConfig(
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    level=logging.INFO,
)
logging.getLogger("httpx").setLevel(logging.WARNING)
logger = logging.getLogger(__name__)

TOKEN = os.environ["TELEGRAM_BOT_TOKEN"]
ALLOWED_USERS_FILE = os.environ.get("ALLOWED_USERS_FILE", "/app/allowed_users.txt")
# Comma/space separated IDs. Lets a deployment work with no bind-mounted file.
ALLOWED_USERS_ENV = os.environ.get("ALLOWED_USERS", "")
WALLETS_BASE  = "/root/.electrum/wallets"
SETTINGS_FILE = "/root/.electrum/bot_settings.json"

# Fee lookup. In Tor mode the entrypoint points these at the mempool.space onion
# service and Tor's SOCKS port, so the fee check never touches the clear net.
MEMPOOL_API_URL = os.environ.get(
    "MEMPOOL_API_URL", "https://mempool.space/api/v1/fees/recommended"
)
TOR_SOCKS = os.environ.get("TOR_SOCKS", "")


# ── Translations ──────────────────────────────────────────────────────────────

STRINGS: dict[str, dict[str, str]] = {
    "en": {
        "welcome": (
            "Welcome to Electrum Wallet Bot! 🔐\n\n"
            "This bot lets you manage Bitcoin watch-only wallets using Electrum on your own server. "
            "Private keys never leave your offline device.\n\n"
            "What you can do:\n"
            "• Restore a wallet from its master public key (xpub)\n"
            "• Check balance in BTC, mBTC, or sats\n"
            "• Get an unused receiving address\n"
            "• Load / close wallets in the Electrum daemon\n"
            "• Check network fee recommendations\n"
            "• Create unsigned transactions for air-gapped signing\n"
            "• Broadcast signed transactions to the network\n\n"
            "Use the buttons below to get started."
        ),
        "not_authorized":     "You are not authorized to use this bot.",
        "invalid_wallet_name": ("❌ Invalid wallet name.\n\n"
                                "Use 1–64 characters: letters, numbers, spaces, dot, dash or underscore. "
                                "It must start with a letter or number.\n\nTry again:"),
        "wallet_exists":      "❌ You already have a wallet named '{name}'. Choose another name:",
        "xpub_received":      "✅ Master public key received — your message was deleted for privacy.",
        "cancelled":          "Cancelled.",
        "cancel_for_button":  "Previous operation cancelled. Tap the button again.",
        "your_wallets":       "Your wallets:",
        "no_wallets":         "No wallets found.",
        "no_wallets_restore": "No wallets found. Use 'Add Wallet' to add one.",
        # Add Wallet
        "add_wallet_intro": (
            "➕ <b>Add Online (Watch-Only) Wallet</b>\n\n"
            "This adds a <b>watch-only</b> wallet using your master public key (xpub). "
            "Your private keys <b>never leave your offline device</b> — the bot can only "
            "view balances, generate addresses, and create <i>unsigned</i> transactions "
            "for you to sign offline.\n\n"
            "You will need:\n"
            "• A wallet name\n"
            "• Your master public key (xpub / zpub / ypub)\n"
            "• A password to encrypt the wallet file\n\n"
            "Let's start — enter a name for this wallet:"
        ),
        "enter_password":   "🔑 Set a password for this wallet:\n\n"
                            "Requirements: at least 8 characters, one uppercase letter, one lowercase letter, and one number.",
        "password_invalid": "❌ Password does not meet requirements.\n\n"
                            "It must be at least 8 characters and contain at least one uppercase letter, "
                            "one lowercase letter, and one number.\n\nTry again:",
        "enter_xpub":       "Enter master public key (xpub / zpub / ypub):",
        "enter_name":       "Enter wallet name:",
        "restoring":        "Restoring '{name}'...",
        "wallet_restored":  "✅ Wallet '{name}' added successfully. Use Load to activate it.",
        # Wallet not loaded
        "wallet_not_loaded": "⚠️ No wallet is loaded.\n\nPlease use <b>Load</b> to load a wallet first.",
        # Load / Close
        "select_wallet":          "Select wallet:",
        "wallet_enter_password":  "Wallet: {name}\nEnter password:",
        "wallet_loaded":          "Wallet '{name}' loaded and selected successfully.",
        "wallet_closed":          "Wallet '{name}' closed successfully.",
        # Balance
        "balance_header":  "Balance ({unit})",
        "confirmed":       "Confirmed:",
        "unconfirmed":     "Unconfirmed:",
        "unmatured":       "Unmatured:",
        # Send
        "send_menu":           "Send menu:",
        "fee_source":          "Select fee source:",
        "recommended_fee":     "Recommended fee: {tip} ({desc})",
        "enter_feerate":       "Enter fee rate (sat/vB):",
        "enter_addr":          "Enter recipient address:",
        "enter_amount":        "Enter amount ({unit}):",
        "invalid_amount":      "Invalid amount: {error}\nTry again:",
        "creating_tx":         "Creating unsigned transaction...",
        "unsigned_tx_header":  "🔐 <b>Unsigned Transaction</b>",
        "signing_instructions": (
            "Sign this with your offline wallet, then use Broadcast.\n\n"
            "📱 Android (Electrum):\n"
            "Send → Paste from clipboard\n\n"
            "💻 Computer (Electrum):\n"
            "Save the file → Send → wrench & screwdriver icon → Read text from file"
        ),
        "enter_signed_tx":    "Paste your signed transaction:",
        "broadcast_success":  "✅ Broadcast successful!\n\nTXID:",
        # Send submenu buttons
        "btn_check_fee":  "Check Mempool Fee",
        "btn_send_tx":    "Send",
        "btn_broadcast":  "Broadcast",
        "btn_send_max":   "Send Max (all)",
        # Fee display
        "fee_mempool_msg": (
            "mempool.space fees:\n"
            "⚡ Fastest:  {fastest} sat/vB\n"
            "🕐 30 min:   {half_hour} sat/vB\n"
            "🕑 1 hour:   {hour} sat/vB\n"
            "💰 Economy:  {economy} sat/vB\n"
            "📉 Minimum:  {minimum} sat/vB"
        ),
        # Language
        "language_select":   "Select language:",
        "language_changed":  "Language changed to English 🇬🇧",
        # Generic
        "error":       "Error:\n{error}",
        "unknown_op":  "Unknown operation.",
        "done":        "Done.",
    },
    "fa": {
        "welcome": (
            "به ربات کیف پول الکترام خوش آمدید! 🔐\n\n"
            "این ربات مدیریت کیف پول‌های بیتکوین را از طریق الکترام روی سرور شخصی شما فراهم می‌کند. "
            "کلیدهای خصوصی هرگز دستگاه آفلاین شما را ترک نمی‌کنند.\n\n"
            "امکانات:\n"
            "• بازیابی کیف پول از کلید عمومی اصلی (xpub)\n"
            "• مشاهده موجودی به BTC، mBTC یا ساتوشی\n"
            "• دریافت آدرس استفاده نشده\n"
            "• بارگذاری و بستن کیف پول در Electrum daemon\n"
            "• بررسی کارمزد شبکه\n"
            "• ایجاد تراکنش بدون امضا برای امضای آفلاین\n"
            "• ارسال تراکنش امضا شده به شبکه\n\n"
            "از دکمه‌های زیر شروع کنید."
        ),
        "not_authorized":     "شما مجاز به استفاده از این ربات نیستید.",
        "invalid_wallet_name": ("❌ نام کیف‌پول نامعتبر است.\n\n"
                                "از ۱ تا ۶۴ کاراکتر استفاده کنید: حروف، اعداد، فاصله، نقطه، خط تیره یا زیرخط. "
                                "باید با حرف یا عدد شروع شود.\n\nدوباره تلاش کنید:"),
        "wallet_exists":      "❌ کیف‌پولی با نام '{name}' از قبل دارید. نام دیگری انتخاب کنید:",
        "xpub_received":      "✅ کلید عمومی اصلی دریافت شد — پیام شما برای حفظ حریم خصوصی حذف شد.",
        "cancelled":          "لغو شد.",
        "cancel_for_button":  "عملیات قبلی لغو شد. دوباره دکمه را فشار دهید.",
        "your_wallets":       "کیف پول‌های شما:",
        "no_wallets":         "کیف پولی یافت نشد.",
        "no_wallets_restore": "کیف پولی یافت نشد. از 'Add Wallet' برای افزودن استفاده کنید.",
        # Add Wallet
        "add_wallet_intro": (
            "➕ <b>افزودن کیف پول آنلاین (فقط مشاهده)</b>\n\n"
            "این کیف پول با استفاده از کلید عمومی اصلی (xpub) شما ایجاد می‌شود. "
            "کلیدهای خصوصی <b>هرگز دستگاه آفلاین شما را ترک نمی‌کنند</b> — ربات فقط می‌تواند "
            "موجودی را مشاهده کند، آدرس تولید کند و تراکنش‌های <i>بدون امضا</i> ایجاد کند.\n\n"
            "به موارد زیر نیاز دارید:\n"
            "• نام کیف پول\n"
            "• کلید عمومی اصلی (xpub / zpub / ypub)\n"
            "• رمز عبور برای رمزگذاری فایل کیف پول\n\n"
            "بیایید شروع کنیم — نامی برای این کیف پول وارد کنید:"
        ),
        "enter_password":   "🔑 رمز عبور این کیف پول را تعیین کنید:\n\n"
                            "الزامات: حداقل ۸ کاراکتر، یک حرف بزرگ، یک حرف کوچک و یک عدد.",
        "password_invalid": "❌ رمز عبور الزامات را برآورده نمی‌کند.\n\n"
                            "باید حداقل ۸ کاراکتر داشته باشد و شامل حداقل یک حرف بزرگ، "
                            "یک حرف کوچک و یک عدد باشد.\n\nدوباره امتحان کنید:",
        "enter_xpub":       "کلید عمومی اصلی (xpub / zpub / ypub) را وارد کنید:",
        "enter_name":       "نام کیف پول را وارد کنید:",
        "restoring":        "در حال افزودن '{name}'...",
        "wallet_restored":  "✅ کیف پول '{name}' با موفقیت افزوده شد. برای فعال‌سازی از Load استفاده کنید.",
        # Wallet not loaded
        "wallet_not_loaded": "⚠️ هیچ کیف پولی بارگذاری نشده است.\n\nلطفاً ابتدا از <b>Load</b> برای بارگذاری کیف پول استفاده کنید.",
        # Load / Close
        "select_wallet":         "کیف پول را انتخاب کنید:",
        "wallet_enter_password": "کیف پول: {name}\nرمز عبور را وارد کنید:",
        "wallet_loaded":         "کیف پول '{name}' با موفقیت بارگذاری و انتخاب شد.",
        "wallet_closed":         "کیف پول '{name}' با موفقیت بسته شد.",
        # Balance
        "balance_header":  "موجودی ({unit})",
        "confirmed":       "تأیید شده:",
        "unconfirmed":     "تأیید نشده:",
        "unmatured":       "بلوغ نیافته:",
        # Send
        "send_menu":          "منوی ارسال:",
        "fee_source":         "منبع کارمزد را انتخاب کنید:",
        "recommended_fee":    "کارمزد پیشنهادی: {tip} ({desc})",
        "enter_feerate":      "نرخ کارمزد را وارد کنید (sat/vB):",
        "enter_addr":         "آدرس گیرنده را وارد کنید:",
        "enter_amount":       "مقدار ({unit}) را وارد کنید:",
        "invalid_amount":     "مقدار نامعتبر: {error}\nدوباره وارد کنید:",
        "creating_tx":        "در حال ایجاد تراکنش بدون امضا...",
        "unsigned_tx_header": "🔐 <b>تراکنش بدون امضا</b>",
        "signing_instructions": (
            "این تراکنش را با کیف پول آفلاین خود امضا کنید، سپس از Broadcast استفاده کنید.\n\n"
            "📱 اندروید (الکترام):\n"
            "Send → Paste from clipboard\n\n"
            "💻 کامپیوتر (الکترام):\n"
            "فایل را ذخیره کنید → Send → آیکون آچار و پیچ‌گوشتی → Read text from file"
        ),
        "enter_signed_tx":   "تراکنش امضا شده را الصاق کنید:",
        "broadcast_success": "✅ تراکنش با موفقیت ارسال شد!\n\nشناسه تراکنش:",
        # Send submenu buttons
        "btn_check_fee":  "بررسی کارمزد",
        "btn_send_tx":    "ارسال",
        "btn_broadcast":  "انتشار",
        "btn_send_max":   "ارسال حداکثر (همه)",
        # Fee display
        "fee_mempool_msg": (
            "کارمزدهای mempool.space:\n"
            "⚡ سریع‌ترین:  {fastest} sat/vB\n"
            "🕐 ۳۰ دقیقه:  {half_hour} sat/vB\n"
            "🕑 ۱ ساعت:    {hour} sat/vB\n"
            "💰 اقتصادی:   {economy} sat/vB\n"
            "📉 حداقل:     {minimum} sat/vB"
        ),
        # Language
        "language_select":  "زبان را انتخاب کنید:",
        "language_changed": "زبان به فارسی تغییر یافت 🇮🇷",
        # Generic
        "error":      "خطا:\n{error}",
        "unknown_op": "عملیات نامشخص.",
        "done":       "انجام شد.",
    },
}


def t(user_id: int, key: str, **kwargs) -> str:
    lang = get_user_lang(user_id)
    text = STRINGS.get(lang, STRINGS["en"]).get(key) or STRINGS["en"].get(key, key)
    try:
        return text.format(**kwargs) if kwargs else text
    except KeyError:
        return text


# ── Settings persistence ──────────────────────────────────────────────────────

def _load_settings() -> dict:
    try:
        with open(SETTINGS_FILE) as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return {}


def _save_settings(settings: dict) -> None:
    # Written via a temp file + atomic rename: a crash half-way through can
    # otherwise leave a truncated settings file, losing every user's wallet
    # selection and language.
    d = os.path.dirname(SETTINGS_FILE)
    os.makedirs(d, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=d, prefix=".bot_settings.", suffix=".tmp")
    try:
        with os.fdopen(fd, "w") as f:
            json.dump(settings, f)
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp, SETTINGS_FILE)
    except BaseException:
        try:
            os.unlink(tmp)
        except OSError:
            pass
        raise


def get_user_unit(user_id: int) -> str:
    return _load_settings().get(str(user_id), {}).get("unit", "BTC")


def set_user_unit(user_id: int, unit: str) -> None:
    s = _load_settings()
    s.setdefault(str(user_id), {})["unit"] = unit
    _save_settings(s)


def get_user_lang(user_id: int) -> str:
    return _load_settings().get(str(user_id), {}).get("lang", "en")


def set_user_lang(user_id: int, lang: str) -> None:
    s = _load_settings()
    s.setdefault(str(user_id), {})["lang"] = lang
    _save_settings(s)


# ── Unit conversion ───────────────────────────────────────────────────────────

def btc_to_unit(btc_str: str, unit: str) -> str:
    try:
        btc = Decimal(btc_str)
    except InvalidOperation:
        return f"{btc_str} {unit}"
    if unit == "BTC":
        return f"{btc:.8f} BTC"
    if unit == "mBTC":
        return f"{btc * 1000:.5f} mBTC"
    return f"{int(btc * Decimal('100000000')):,} sats"


def unit_to_btc(amount_str: str, unit: str) -> str:
    try:
        amount = Decimal(amount_str.replace(",", ""))
    except InvalidOperation:
        raise ValueError(amount_str)
    if unit == "BTC":
        return f"{amount:.8f}"
    if unit == "mBTC":
        return f"{amount / Decimal('1000'):.8f}"
    return f"{amount / Decimal('100000000'):.8f}"


def validate_password(password: str) -> bool:
    if len(password) < 8:
        return False
    if not any(c.isupper() for c in password):
        return False
    if not any(c.islower() for c in password):
        return False
    if not any(c.isdigit() for c in password):
        return False
    return True


def get_loaded_wallet(user_id: int) -> str | None:
    return _load_settings().get(str(user_id), {}).get("loaded_wallet")


def set_loaded_wallet(user_id: int, name: str | None) -> None:
    s = _load_settings()
    s.setdefault(str(user_id), {})["loaded_wallet"] = name
    _save_settings(s)


def format_balance(data: dict, unit: str, lang: str) -> str:
    s = STRINGS.get(lang, STRINGS["en"])
    confirmed   = data.get("confirmed",   "0") or "0"
    unconfirmed = data.get("unconfirmed", "0") or "0"
    unmatured   = data.get("unmatured",   "0") or "0"
    lines = [s["balance_header"].format(unit=unit)]
    lines.append(f"{s['confirmed']}  {btc_to_unit(confirmed, unit)}")
    if Decimal(unconfirmed) != 0:
        lines.append(f"{s['unconfirmed']}  {btc_to_unit(unconfirmed, unit)}")
    if Decimal(unmatured) != 0:
        lines.append(f"{s['unmatured']}  {btc_to_unit(unmatured, unit)}")
    return "\n".join(lines)


# ── Keyboards ─────────────────────────────────────────────────────────────────

MAIN_KEYBOARD = ReplyKeyboardMarkup(
    [["Add Wallet", "List"], ["Balance", "Receive"], ["Load", "Close"], ["Send"], ["🌐 Language"]],
    resize_keyboard=True,
    is_persistent=True,
)

def submenu_kb(user_id: int) -> InlineKeyboardMarkup:
    return InlineKeyboardMarkup([
        [InlineKeyboardButton(t(user_id, "btn_check_fee"),  callback_data="check_fee")],
        [InlineKeyboardButton(t(user_id, "btn_send_tx"),    callback_data="do_send")],
        [InlineKeyboardButton(t(user_id, "btn_broadcast"),  callback_data="do_broadcast")],
    ])


def fee_kb(user_id: int) -> InlineKeyboardMarkup:
    return InlineKeyboardMarkup([[
        InlineKeyboardButton("Electrum",      callback_data="fee_electrum"),
        InlineKeyboardButton("mempool.space", callback_data="fee_mempool"),
    ]])


def amount_kb(user_id: int) -> InlineKeyboardMarkup:
    return InlineKeyboardMarkup([[
        InlineKeyboardButton(t(user_id, "btn_send_max"), callback_data="send_max"),
    ]])

LANG_KB = InlineKeyboardMarkup([[
    InlineKeyboardButton("🇬🇧 English", callback_data="lang:en"),
    InlineKeyboardButton("🇮🇷 فارسی",  callback_data="lang:fa"),
]])


def unit_keyboard(current: str) -> InlineKeyboardMarkup:
    def btn(u: str) -> InlineKeyboardButton:
        return InlineKeyboardButton(f"● {u}" if u == current else u, callback_data=f"unit:{u}")
    return InlineKeyboardMarkup([[btn("BTC"), btn("mBTC"), btn("sats")]])


MENU_FILTER = filters.Regex(r"^(Add Wallet|List|Balance|Receive|Load|Close|Send|🌐 Language)$")
NOT_MENU    = ~MENU_FILTER

# ── Conversation states ───────────────────────────────────────────────────────
RESTORE_NAME, RESTORE_XPUB, RESTORE_PASSWORD = range(3)
OP_SELECT, OP_PASSWORD = range(2)
S_MENU, S_FEE_SRC, S_RATE, S_PASS, S_ADDR, S_AMT, S_BCAST = range(7)


# ── Core helpers ──────────────────────────────────────────────────────────────

def load_allowed_users() -> set[int]:
    users: set[int] = set()
    for chunk in re.split(r"[,\s]+", ALLOWED_USERS_ENV):
        chunk = chunk.strip()
        if chunk.isdigit():
            users.add(int(chunk))
    try:
        with open(ALLOWED_USERS_FILE) as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                uid = line.split(",", 1)[0].strip()
                if uid.isdigit():
                    users.add(int(uid))
    except (FileNotFoundError, IsADirectoryError, PermissionError):
        # Not fatal — ALLOWED_USERS may be supplying the list instead
        logger.warning("Allowed users file unreadable: %s", ALLOWED_USERS_FILE)
    if not users:
        logger.warning("No allowed users configured — the bot will reject everyone")
    return users


def user_wallet_path(user_id: int) -> str:
    path = f"{WALLETS_BASE}/{user_id}"
    os.makedirs(path, exist_ok=True)
    return path


# Wallet names become file names. Without this a name like "../<other-id>/x"
# would read or overwrite another user's wallet, or clobber bot_settings.json.
WALLET_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9 ._-]{0,63}$")


def is_valid_wallet_name(name: str) -> bool:
    return bool(name) and bool(WALLET_NAME_RE.fullmatch(name)) and "/" not in name


def wallet_file(user_id: int, name: str) -> str | None:
    """Absolute path of a wallet, or None if the name escapes the user's dir."""
    if not is_valid_wallet_name(name):
        return None
    base = os.path.realpath(user_wallet_path(user_id))
    path = os.path.realpath(os.path.join(base, name))
    # Belt and braces: even a name that passed the regex must resolve inside base
    if path != base and not path.startswith(base + os.sep):
        logger.warning("Rejected wallet path outside user dir: user=%s name=%r", user_id, name)
        return None
    return path


def list_user_wallets(user_id: int) -> list[str]:
    return sorted(
        w for w in os.listdir(user_wallet_path(user_id)) if is_valid_wallet_name(w)
    )


def run_electrum(*args: str) -> tuple[str, str, int]:
    result = subprocess.run(
        ["electrum", *args], capture_output=True, text=True, timeout=30,
    )
    return result.stdout.strip(), result.stderr.strip(), result.returncode


async def auth_check(update: Update) -> bool:
    """Gate every entry point. Also called from callback handlers, where
    update.message is None, so the reply has to go through the query."""
    uid = update.effective_user.id
    if uid in load_allowed_users():
        return True
    if update.callback_query:
        await update.callback_query.answer(t(uid, "not_authorized"), show_alert=True)
    elif update.message:
        await update.message.reply_text(t(uid, "not_authorized"))
    return False


async def _delete_quietly(message) -> None:
    """Best-effort removal of a message holding a secret. Telegram refuses for
    messages older than 48h, and the secret is still worth scrubbing when it
    works, so a failure here must not abort the flow."""
    try:
        await message.delete()
    except Exception as e:  # noqa: BLE001 - never let cleanup break the conversation
        logger.warning("Could not delete sensitive message: %s", e)


async def reply(update: Update, text: str, **kwargs) -> None:
    if update.callback_query:
        await update.callback_query.edit_message_text(text, **kwargs)
    else:
        await update.message.reply_text(text, reply_markup=MAIN_KEYBOARD, **kwargs)


async def cancel(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    uid = update.effective_user.id
    context.user_data.clear()
    await update.message.reply_text(t(uid, "cancelled"), reply_markup=MAIN_KEYBOARD)
    return ConversationHandler.END


async def cancel_for_button(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    uid = update.effective_user.id
    context.user_data.clear()
    await update.message.reply_text(t(uid, "cancel_for_button"), reply_markup=MAIN_KEYBOARD)
    return ConversationHandler.END


CONV_FALLBACKS = [
    CommandHandler("cancel", cancel),
    MessageHandler(MENU_FILTER, cancel_for_button),
]


# ── Simple commands ───────────────────────────────────────────────────────────

async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if not await auth_check(update):
        return
    uid = update.effective_user.id
    await update.message.reply_text(t(uid, "welcome"), reply_markup=MAIN_KEYBOARD)


async def list_wallets(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if not await auth_check(update):
        return
    uid     = update.effective_user.id
    wallets = list_user_wallets(uid)
    if wallets:
        msg = t(uid, "your_wallets") + "\n" + "\n".join(wallets)
    else:
        msg = t(uid, "no_wallets")
    await update.message.reply_text(msg, reply_markup=MAIN_KEYBOARD)


async def balance_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if not await auth_check(update):
        return
    uid  = update.effective_user.id
    if not get_loaded_wallet(uid):
        await update.message.reply_text(t(uid, "wallet_not_loaded"), parse_mode=ParseMode.HTML, reply_markup=MAIN_KEYBOARD)
        return
    unit = get_user_unit(uid)
    lang = get_user_lang(uid)
    wallet_path = wallet_file(uid, get_loaded_wallet(uid))
    if wallet_path is None:
        await update.message.reply_text(t(uid, "wallet_not_loaded"), parse_mode=ParseMode.HTML, reply_markup=MAIN_KEYBOARD)
        return
    stdout, stderr, rc = run_electrum("getbalance", "-w", wallet_path)
    if rc != 0:
        await update.message.reply_text(t(uid, "error", error=stderr or stdout), reply_markup=MAIN_KEYBOARD)
        return
    try:
        data = json.loads(stdout)
        msg  = format_balance(data, unit, lang)
    except (json.JSONDecodeError, InvalidOperation):
        await update.message.reply_text(stdout, reply_markup=MAIN_KEYBOARD)
        return
    await update.message.reply_text(msg, reply_markup=unit_keyboard(unit))


async def unit_switch(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if not await auth_check(update):
        return
    query = update.callback_query
    await query.answer()
    uid  = query.from_user.id
    unit = query.data.split(":", 1)[1]
    set_user_unit(uid, unit)
    lang = get_user_lang(uid)
    loaded = get_loaded_wallet(uid)
    if not loaded:
        await query.edit_message_text(t(uid, "wallet_not_loaded"), parse_mode=ParseMode.HTML)
        return
    wallet_path = wallet_file(uid, loaded)
    if wallet_path is None:
        await query.edit_message_text(t(uid, "wallet_not_loaded"), parse_mode=ParseMode.HTML)
        return
    stdout, stderr, rc = run_electrum("getbalance", "-w", wallet_path)
    if rc != 0:
        await query.edit_message_text(t(uid, "error", error=stderr or stdout))
        return
    try:
        data = json.loads(stdout)
        msg  = format_balance(data, unit, lang)
    except (json.JSONDecodeError, InvalidOperation):
        msg = stdout
    await query.edit_message_text(msg, reply_markup=unit_keyboard(unit))


async def receive_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if not await auth_check(update):
        return
    uid = update.effective_user.id
    loaded = get_loaded_wallet(uid)
    if not loaded:
        await update.message.reply_text(t(uid, "wallet_not_loaded"), parse_mode=ParseMode.HTML, reply_markup=MAIN_KEYBOARD)
        return
    wallet_path = wallet_file(uid, loaded)
    if wallet_path is None:
        await update.message.reply_text(t(uid, "wallet_not_loaded"), parse_mode=ParseMode.HTML, reply_markup=MAIN_KEYBOARD)
        return
    stdout, stderr, rc = run_electrum("getunusedaddress", "-w", wallet_path)
    msg = stdout if rc == 0 else t(uid, "error", error=stderr or stdout)
    await update.message.reply_text(msg, reply_markup=MAIN_KEYBOARD)


async def lang_menu(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if not await auth_check(update):
        return
    uid = update.effective_user.id
    await update.message.reply_text(t(uid, "language_select"), reply_markup=LANG_KB)


async def lang_switch(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if not await auth_check(update):
        return
    query = update.callback_query
    await query.answer()
    uid  = query.from_user.id
    lang = query.data.split(":", 1)[1]
    set_user_lang(uid, lang)
    await query.edit_message_text(t(uid, "language_changed"))


# ── Restore conversation ──────────────────────────────────────────────────────

async def restore_start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    if not await auth_check(update):
        return ConversationHandler.END
    uid = update.effective_user.id
    await update.message.reply_text(t(uid, "add_wallet_intro"), parse_mode=ParseMode.HTML)
    return RESTORE_NAME


async def restore_got_name(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    uid  = update.effective_user.id
    name = update.message.text.strip()
    if not is_valid_wallet_name(name):
        await update.message.reply_text(t(uid, "invalid_wallet_name"))
        return RESTORE_NAME
    if name in list_user_wallets(uid):
        await update.message.reply_text(t(uid, "wallet_exists", name=name))
        return RESTORE_NAME
    context.user_data["walletname"] = name
    await update.message.reply_text(t(uid, "enter_xpub"))
    return RESTORE_XPUB


async def restore_got_pubkey(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    uid = update.effective_user.id
    context.user_data["pubkey"] = update.message.text.strip()
    # The xpub reveals every address and the full balance history of the wallet.
    # Don't leave it sitting in the Telegram chat forever.
    await _delete_quietly(update.message)
    await context.bot.send_message(update.effective_chat.id, t(uid, "xpub_received"))
    await context.bot.send_message(update.effective_chat.id, t(uid, "enter_password"))
    return RESTORE_PASSWORD


async def restore_got_password(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    uid      = update.effective_user.id
    password = update.message.text.strip()
    await _delete_quietly(update.message)
    if not validate_password(password):
        await context.bot.send_message(update.effective_chat.id, t(uid, "password_invalid"))
        return RESTORE_PASSWORD
    context.user_data["password"] = password
    walletname  = context.user_data["walletname"]
    wallet_path = wallet_file(uid, walletname)
    if wallet_path is None:
        await context.bot.send_message(update.effective_chat.id, t(uid, "invalid_wallet_name"),
                                       reply_markup=MAIN_KEYBOARD)
        context.user_data.clear()
        return ConversationHandler.END
    # The originating message is gone, so send rather than reply
    await context.bot.send_message(update.effective_chat.id, t(uid, "restoring", name=walletname))
    stdout, stderr, rc = run_electrum(
        "restore", "--password", context.user_data["password"],
        "--encrypt_file", "true", "-w", wallet_path, context.user_data["pubkey"],
    )
    msg = t(uid, "wallet_restored", name=walletname) if rc == 0 else t(uid, "error", error=stderr or stdout)
    await context.bot.send_message(update.effective_chat.id, msg, reply_markup=MAIN_KEYBOARD)
    context.user_data.clear()
    return ConversationHandler.END


# ── Load / Close conversation ─────────────────────────────────────────────────

async def op_start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    if not await auth_check(update):
        return ConversationHandler.END
    uid = update.effective_user.id
    raw = (update.message.text or "").lstrip("/").split()[0].lower()
    context.user_data["op"] = raw
    wallets = list_user_wallets(uid)
    if not wallets:
        await update.message.reply_text(t(uid, "no_wallets"), reply_markup=MAIN_KEYBOARD)
        return ConversationHandler.END
    if len(wallets) == 1:
        return await _wallet_selected(update, context, wallets[0])
    keyboard = [[InlineKeyboardButton(w, callback_data=w)] for w in wallets]
    await update.message.reply_text(t(uid, "select_wallet"), reply_markup=InlineKeyboardMarkup(keyboard))
    return OP_SELECT


async def op_wallet_selected(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    await update.callback_query.answer()
    uid  = update.effective_user.id
    name = update.callback_query.data
    # callback_data is echoed back by the client — confirm it names a real wallet
    if name not in list_user_wallets(uid):
        await update.callback_query.edit_message_text(t(uid, "invalid_wallet_name"))
        return ConversationHandler.END
    return await _wallet_selected(update, context, name)


async def _wallet_selected(
    update: Update, context: ContextTypes.DEFAULT_TYPE, walletname: str
) -> int:
    uid = update.effective_user.id
    context.user_data["walletname"] = walletname
    if context.user_data.get("op") == "load":
        await reply(update, t(uid, "wallet_enter_password", name=walletname))
        return OP_PASSWORD
    return await _execute_op(update, context)


async def op_got_password(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    context.user_data["password"] = update.message.text.strip()
    await _delete_quietly(update.message)
    return await _execute_op(update, context)


async def _execute_op(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    uid        = update.effective_user.id
    op         = context.user_data.get("op", "")
    walletname = context.user_data.get("walletname", "")
    wallet_path = wallet_file(uid, walletname)
    if wallet_path is None:
        await reply(update, t(uid, "invalid_wallet_name"))
        context.user_data.clear()
        return ConversationHandler.END
    cmd_map = {
        "load":  ("load_wallet",  "-w", wallet_path, "--password", context.user_data.get("password", "")),
        "close": ("close_wallet", "-w", wallet_path),
    }
    electrum_cmd = cmd_map.get(op)
    if electrum_cmd:
        _, stderr, rc = run_electrum(*electrum_cmd)
        if rc == 0:
            if op == "load":
                set_loaded_wallet(uid, walletname)
                msg = t(uid, "wallet_loaded", name=walletname)
            else:
                set_loaded_wallet(uid, None)
                msg = t(uid, "wallet_closed", name=walletname)
        else:
            msg = t(uid, "error", error=stderr)
    else:
        msg = t(uid, "unknown_op")
    await reply(update, msg)
    context.user_data.clear()
    return ConversationHandler.END


# ── Send conversation ─────────────────────────────────────────────────────────

async def send_menu(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    if not await auth_check(update):
        return ConversationHandler.END
    uid = update.effective_user.id
    await update.message.reply_text(t(uid, "send_menu"), reply_markup=submenu_kb(uid))
    return S_MENU


async def _show_submenu(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    uid = update.effective_user.id
    await context.bot.send_message(update.effective_chat.id, t(uid, "send_menu"), reply_markup=submenu_kb(uid))
    return S_MENU


async def fee_choice(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    uid = update.callback_query.from_user.id
    await update.callback_query.answer()
    await update.callback_query.edit_message_text(t(uid, "fee_source"), reply_markup=fee_kb(uid))
    return S_FEE_SRC


async def fee_electrum(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    query = update.callback_query
    uid   = query.from_user.id
    await query.answer()
    stdout, stderr, rc = run_electrum("getfeerate")
    if rc == 0:
        try:
            d   = json.loads(stdout)
            msg = t(uid, "recommended_fee", tip=d["tooltip"], desc=d["description"])
        except (json.JSONDecodeError, KeyError):
            msg = stdout
    else:
        msg = t(uid, "error", error=stderr or stdout)
    await query.edit_message_text(msg)
    return await _show_submenu(update, context)


async def fee_mempool(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    query = update.callback_query
    await query.answer()
    uid = query.from_user.id
    try:
        cmd = ["curl", "-sSL", "--max-time", "15"]
        if TOR_SOCKS:
            # --socks5-hostname keeps DNS inside Tor too, so the lookup never
            # leaks either. Set by the Tor entrypoint.
            cmd += ["--socks5-hostname", TOR_SOCKS]
        cmd.append(MEMPOOL_API_URL)
        res = subprocess.run(cmd, capture_output=True, text=True, timeout=20)
        if res.returncode != 0:
            raise RuntimeError(res.stderr.strip() or "fee lookup failed")
        d   = json.loads(res.stdout)
        msg = t(uid, "fee_mempool_msg",
                fastest=d["fastestFee"], half_hour=d["halfHourFee"],
                hour=d["hourFee"], economy=d["economyFee"], minimum=d["minimumFee"])
    except Exception as e:
        msg = t(uid, "error", error=str(e))
    await query.edit_message_text(msg)
    return await _show_submenu(update, context)


async def send_ask_feerate(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    uid = update.callback_query.from_user.id
    await update.callback_query.answer()
    if not get_loaded_wallet(uid):
        await update.callback_query.edit_message_text(t(uid, "wallet_not_loaded"), parse_mode=ParseMode.HTML)
        return ConversationHandler.END
    await update.callback_query.edit_message_text(t(uid, "enter_feerate"))
    return S_RATE


async def send_got_feerate(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    context.user_data["feerate"] = update.message.text.strip()
    uid = update.effective_user.id
    await update.message.reply_text(t(uid, "enter_password"))
    return S_PASS


async def send_got_password(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    context.user_data["password"] = update.message.text.strip()
    uid = update.effective_user.id
    await _delete_quietly(update.message)
    await context.bot.send_message(update.effective_chat.id, t(uid, "enter_addr"))
    return S_ADDR


async def send_got_address(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    context.user_data["address"] = update.message.text.strip()
    uid  = update.effective_user.id
    unit = get_user_unit(uid)
    await update.message.reply_text(t(uid, "enter_amount", unit=unit), reply_markup=amount_kb(uid))
    return S_AMT


async def send_max_cb(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    await update.callback_query.answer()
    await update.callback_query.edit_message_reply_markup(reply_markup=None)
    return await _build_tx(update, context, "!")


async def send_got_amount(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    uid        = update.effective_user.id
    unit       = get_user_unit(uid)
    amount_raw = update.message.text.strip()
    try:
        btc_amount = unit_to_btc(amount_raw, unit)
    except ValueError as e:
        await update.message.reply_text(t(uid, "invalid_amount", error=str(e)))
        return S_AMT
    return await _build_tx(update, context, btc_amount)


async def _build_tx(update: Update, context: ContextTypes.DEFAULT_TYPE, btc_amount: str) -> int:
    uid = update.effective_user.id

    notify = update.callback_query.message if update.callback_query else update.message
    await notify.reply_text(t(uid, "creating_tx"))

    wallet_path = wallet_file(uid, get_loaded_wallet(uid))
    if wallet_path is None:
        await notify.reply_text(t(uid, "wallet_not_loaded"), parse_mode=ParseMode.HTML, reply_markup=MAIN_KEYBOARD)
        context.user_data.clear()
        return ConversationHandler.END
    stdout, stderr, rc = run_electrum(
        "payto",
        "-w", wallet_path,
        "--feerate", context.user_data["feerate"],
        "--rbf", "true",
        "--password", context.user_data["password"],
        context.user_data["address"],
        btc_amount,
    )

    if rc != 0:
        await notify.reply_text(t(uid, "error", error=stderr or stdout), reply_markup=MAIN_KEYBOARD)
        context.user_data.clear()
        return ConversationHandler.END

    psbt = stdout.strip()
    await notify.reply_text(
        f"{t(uid, 'unsigned_tx_header')}\n\n<pre>{psbt}</pre>",
        parse_mode=ParseMode.HTML,
    )

    file_obj = io.BytesIO(psbt.encode())
    await context.bot.send_document(
        chat_id=update.effective_chat.id,
        document=file_obj,
        filename="unsigned_tx.txt",
        caption="Unsigned transaction",
    )

    await notify.reply_text(t(uid, "signing_instructions"), reply_markup=MAIN_KEYBOARD)
    context.user_data.clear()
    return ConversationHandler.END


async def broadcast_ask(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    uid = update.callback_query.from_user.id
    await update.callback_query.answer()
    await update.callback_query.edit_message_text(t(uid, "enter_signed_tx"))
    return S_BCAST


async def broadcast_got_tx(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    uid = update.effective_user.id
    stdout, stderr, rc = run_electrum("broadcast", update.message.text.strip())
    if rc == 0:
        await update.message.reply_text(
            f"{t(uid, 'broadcast_success')}\n<code>{stdout}</code>",
            parse_mode=ParseMode.HTML,
            reply_markup=MAIN_KEYBOARD,
        )
    else:
        await update.message.reply_text(t(uid, "error", error=stderr or stdout), reply_markup=MAIN_KEYBOARD)
    return ConversationHandler.END


# ── App wiring ────────────────────────────────────────────────────────────────

def main() -> None:
    app = Application.builder().token(TOKEN).build()

    for cmd in ("start", "help"):
        app.add_handler(CommandHandler(cmd, start))
    app.add_handler(CommandHandler("list",    list_wallets))
    app.add_handler(CommandHandler("balance", balance_cmd))
    app.add_handler(CommandHandler("receive", receive_cmd))

    app.add_handler(MessageHandler(filters.Regex(r"^List$"),         list_wallets))
    app.add_handler(MessageHandler(filters.Regex(r"^Balance$"),      balance_cmd))
    app.add_handler(MessageHandler(filters.Regex(r"^Receive$"),      receive_cmd))
    app.add_handler(MessageHandler(filters.Regex(r"^🌐 Language$"),  lang_menu))

    # Global inline callbacks (registered before ConversationHandlers — patterns are unique)
    app.add_handler(CallbackQueryHandler(unit_switch, pattern=r"^unit:"))
    app.add_handler(CallbackQueryHandler(lang_switch, pattern=r"^lang:"))

    app.add_handler(ConversationHandler(
        entry_points=[
            CommandHandler("restore", restore_start),
            MessageHandler(filters.Regex(r"^Add Wallet$"), restore_start),
        ],
        states={
            RESTORE_NAME:     [MessageHandler(filters.TEXT & ~filters.COMMAND & NOT_MENU, restore_got_name)],
            RESTORE_XPUB:     [MessageHandler(filters.TEXT & ~filters.COMMAND & NOT_MENU, restore_got_pubkey)],
            RESTORE_PASSWORD: [MessageHandler(filters.TEXT & ~filters.COMMAND & NOT_MENU, restore_got_password)],
        },
        fallbacks=CONV_FALLBACKS,
    ))

    app.add_handler(ConversationHandler(
        entry_points=[
            CommandHandler("load",  op_start),
            CommandHandler("close", op_start),
            MessageHandler(filters.Regex(r"^(Load|Close)$"), op_start),
        ],
        states={
            OP_SELECT:   [CallbackQueryHandler(op_wallet_selected)],
            OP_PASSWORD: [MessageHandler(filters.TEXT & ~filters.COMMAND & NOT_MENU, op_got_password)],
        },
        fallbacks=CONV_FALLBACKS,
    ))

    app.add_handler(ConversationHandler(
        entry_points=[
            CommandHandler("send", send_menu),
            MessageHandler(filters.Regex(r"^Send$"), send_menu),
        ],
        states={
            S_MENU: [
                CallbackQueryHandler(fee_choice,       pattern="^check_fee$"),
                CallbackQueryHandler(send_ask_feerate, pattern="^do_send$"),
                CallbackQueryHandler(broadcast_ask,    pattern="^do_broadcast$"),
            ],
            S_FEE_SRC: [
                CallbackQueryHandler(fee_electrum, pattern="^fee_electrum$"),
                CallbackQueryHandler(fee_mempool,  pattern="^fee_mempool$"),
            ],
            S_RATE:  [MessageHandler(filters.TEXT & ~filters.COMMAND & NOT_MENU, send_got_feerate)],
            S_PASS:  [MessageHandler(filters.TEXT & ~filters.COMMAND & NOT_MENU, send_got_password)],
            S_ADDR:  [MessageHandler(filters.TEXT & ~filters.COMMAND & NOT_MENU, send_got_address)],
            S_AMT: [
                CallbackQueryHandler(send_max_cb, pattern="^send_max$"),
                MessageHandler(filters.TEXT & ~filters.COMMAND & NOT_MENU, send_got_amount),
            ],
            S_BCAST: [MessageHandler(filters.TEXT & ~filters.COMMAND & NOT_MENU, broadcast_got_tx)],
        },
        fallbacks=CONV_FALLBACKS,
    ))

    logger.info("Bot starting...")
    app.run_polling(allowed_updates=Update.ALL_TYPES)


if __name__ == "__main__":
    main()