"""Encrypted wallet vault — AES-256-GCM + Argon2id, stored in SQLite.

Trust principle: The vault password NEVER leaves the server. If you use
client-side generation, the private key is encrypted before it ever reaches
the server. The server stores opaque blobs it cannot decrypt.

Server-side generation (enterprise vault) encrypts keys at rest with
the vault password. The password is configurable and never logged.
"""

from __future__ import annotations

import base64
import json
import logging
import secrets
import sqlite3
import threading
import time
import queue
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterator, Optional

from .config import cfg

logger = logging.getLogger("wp.vault")

try:
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    from cryptography.hazmat.primitives.kdf.argon2 import Argon2id
except ImportError:
    raise RuntimeError(
        "cryptography is required. Install it: pip install cryptography>=42.0.0"
    )


@dataclass
class WalletEntry:
    id: str
    chain: str
    address: str
    label: str
    tags: list[str] = field(default_factory=list)
    group: str = ""
    created_at: float = 0.0
    encrypted_key: str = ""
    public_key: str = ""
    derivation_path: str = ""
    hd_path: str = ""
    mnemonic_encrypted: str = ""
    encrypted: bool = False
    balance_usd: float = 0.0
    notes: str = ""


class Vault:
    """Encrypted wallet vault with pooled SQLite connections.

    Connection pool: up to POOL_SIZE reusable connections, thread-safe.
    WAL mode + busy_timeout prevent 'database is locked' under concurrent load.
    """

    POOL_SIZE = 8
    BUSY_TIMEOUT = 5000  # ms before raising locked error

    def __init__(self, db_path: Path):
        self.db_path = db_path
        self._pool: queue.Queue[sqlite3.Connection] = queue.Queue(maxsize=self.POOL_SIZE)
        self._lock = threading.Lock()
        self._total_conns = 0
        if not cfg.vault_password:
            raise RuntimeError(
                "WalletPress vault requires encryption. "
                "Set WP_VAULT_PASSWORD environment variable to a strong password. "
                "Example: export WP_VAULT_PASSWORD='$(openssl rand -base64 32)'"
            )
        self._init_db()

    def _new_conn(self) -> sqlite3.Connection:
        conn = sqlite3.connect(str(self.db_path), timeout=self.BUSY_TIMEOUT / 1000, check_same_thread=False)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA busy_timeout=%d" % self.BUSY_TIMEOUT)
        conn.execute("PRAGMA foreign_keys=ON")
        return conn

    @contextmanager
    def _conn(self) -> Iterator[sqlite3.Connection]:
        """Get a connection from the pool (or create one). Auto-returns on exit."""
        try:
            conn = self._pool.get_nowait()
        except queue.Empty:
            with self._lock:
                if self._total_conns < self.POOL_SIZE:
                    conn = self._new_conn()
                    self._total_conns += 1
                else:
                    conn = self._new_conn()  # overflow — will close after use
        try:
            yield conn
        finally:
            try:
                self._pool.put_nowait(conn)
            except queue.Full:
                conn.close()
                with self._lock:
                    self._total_conns -= 1

    SCHEMA_VERSION = 1

    def _init_db(self):
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        with self._new_conn() as conn:
            conn.executescript("""
            CREATE TABLE IF NOT EXISTS wallets (
                id TEXT PRIMARY KEY,
                chain TEXT NOT NULL,
                address TEXT NOT NULL,
                label TEXT DEFAULT '',
                tags TEXT DEFAULT '[]',
                wallet_group TEXT DEFAULT '',
                created_at REAL NOT NULL,
                encrypted_key TEXT DEFAULT '',
                public_key TEXT DEFAULT '',
                derivation_path TEXT DEFAULT '',
                hd_path TEXT DEFAULT '',
                mnemonic_encrypted TEXT DEFAULT '',
                encrypted INTEGER DEFAULT 0,
                balance_usd REAL DEFAULT 0.0,
                notes TEXT DEFAULT '',
                updated_at REAL DEFAULT 0
            );
            CREATE INDEX IF NOT EXISTS idx_wallets_chain ON wallets(chain);
            CREATE INDEX IF NOT EXISTS idx_wallets_address ON wallets(address);
            CREATE VIRTUAL TABLE IF NOT EXISTS wallets_fts USING fts5(
                id, chain, address, label, notes,
                content='wallets',
                content_rowid='rowid'
            );
            CREATE TABLE IF NOT EXISTS rotations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                wallet_id TEXT NOT NULL,
                new_address TEXT NOT NULL,
                rotated_at REAL NOT NULL,
                reason TEXT DEFAULT '',
                FOREIGN KEY(wallet_id) REFERENCES wallets(id)
            );
            CREATE TABLE IF NOT EXISTS _schema_version (
                version INTEGER PRIMARY KEY,
                applied_at REAL NOT NULL
            );
        """)
            conn.commit()
            self._migrate(conn)

    def _migrate(self, conn):
        # P3-5: Schema migrations are now handled by Alembic. This stub
        # only stamps the schema version table. Kept for backward compat
        # with existing DBs that have _schema_version rows.
        row = conn.execute("SELECT MAX(version) FROM _schema_version").fetchone()
        current = row[0] if row and row[0] else 0
        if current >= self.SCHEMA_VERSION:
            return
        with self._lock:
            conn.execute(
                "INSERT OR IGNORE INTO _schema_version (version, applied_at) VALUES (?, ?)",
                (self.SCHEMA_VERSION, time.time()),
            )
            conn.commit()
            logger.info(
                f"Vault schema version stamped at v{self.SCHEMA_VERSION}. "
                f"For real schema migrations, run `alembic upgrade head`."
            )

    def encrypt_key(self, plaintext: str) -> str:
        salt = secrets.token_bytes(16)
        kdf = Argon2id(salt, 32, 3, 4, 65536)
        key = kdf.derive(cfg.vault_password.encode())
        aesgcm = AESGCM(key)
        nonce = secrets.token_bytes(12)
        ct = aesgcm.encrypt(nonce, plaintext.encode(), None)
        return base64.b64encode(salt + nonce + ct).decode()

    def decrypt_key(self, encrypted: str) -> str:
        data = base64.b64decode(encrypted)
        salt, nonce, ct = data[:16], data[16:28], data[28:]
        kdf = Argon2id(salt, 32, 3, 4, 65536)
        key = kdf.derive(cfg.vault_password.encode())
        aesgcm = AESGCM(key)
        return aesgcm.decrypt(nonce, ct, None).decode()

    def put(self, wallet: WalletEntry) -> bool:
        with self._conn() as conn:
            try:
                try:
                    conn.execute("ALTER TABLE wallets ADD COLUMN wallet_group TEXT DEFAULT ''")
                    conn.commit()
                except Exception:
                    pass

                conn.execute(
                    """INSERT OR REPLACE INTO wallets
                    (id, chain, address, label, tags, wallet_group, created_at, encrypted_key,
                     public_key, derivation_path, hd_path, mnemonic_encrypted,
                     encrypted, balance_usd, notes, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                    (
                        wallet.id, wallet.chain, wallet.address, wallet.label,
                        json.dumps(wallet.tags), wallet.group, wallet.created_at or time.time(),
                        wallet.encrypted_key, wallet.public_key,
                        wallet.derivation_path, wallet.hd_path,
                        wallet.mnemonic_encrypted, int(wallet.encrypted),
                        wallet.balance_usd, wallet.notes, time.time(),
                    ),
                )
                conn.commit()
                self.update_fts(wallet)
                return True
            except Exception as e:
                logger.error(f"Vault put failed: {e}")
                return False

    def get(self, wallet_id: str) -> Optional[WalletEntry]:
        with self._conn() as conn:
            row = conn.execute("SELECT * FROM wallets WHERE id = ?", (wallet_id,)).fetchone()
            if not row:
                return None
            return self._row_to_entry(row)

    def get_by_address(self, address: str) -> Optional[WalletEntry]:
        with self._conn() as conn:
            row = conn.execute("SELECT * FROM wallets WHERE address = ?", (address,)).fetchone()
            if not row:
                return None
            return self._row_to_entry(row)

    def list(self, chain: str = "", limit: int = 100, offset: int = 0) -> list[WalletEntry]:
        with self._conn() as conn:
            if chain:
                rows = conn.execute(
                    "SELECT * FROM wallets WHERE chain = ? ORDER BY created_at DESC LIMIT ? OFFSET ?",
                    (chain, limit, offset),
                ).fetchall()
            else:
                rows = conn.execute(
                    "SELECT * FROM wallets ORDER BY created_at DESC LIMIT ? OFFSET ?",
                    (limit, offset),
                ).fetchall()
            return [self._row_to_entry(r) for r in rows]

    def search(self, query: str) -> list[WalletEntry]:
        with self._conn() as conn:
            try:
                rows = conn.execute(
                    "SELECT wallets.* FROM wallets_fts JOIN wallets ON wallets_fts.id = wallets.id WHERE wallets_fts MATCH ? ORDER BY rank LIMIT 50",
                    (query,),
                ).fetchall()
            except Exception:
                like = f"%{query}%"
                rows = conn.execute(
                    "SELECT * FROM wallets WHERE address LIKE ? OR label LIKE ? OR chain LIKE ? OR notes LIKE ? LIMIT 50",
                    (like, like, like, like),
                ).fetchall()
            return [self._row_to_entry(r) for r in rows]

    def update_fts(self, wallet: WalletEntry):
        with self._conn() as conn:
            try:
                conn.execute(
                    "INSERT OR REPLACE INTO wallets_fts (id, chain, address, label, notes) VALUES (?, ?, ?, ?, ?)",
                    (wallet.id, wallet.chain, wallet.address, wallet.label, wallet.notes),
                )
                conn.commit()
            except Exception:
                pass

    def delete(self, wallet_id: str) -> bool:
        with self._conn() as conn:
            conn.execute("DELETE FROM wallets WHERE id = ?", (wallet_id,))
            conn.commit()
            return conn.total_changes > 0

    def count(self, chain: str = "") -> int:
        with self._conn() as conn:
            if chain:
                return conn.execute("SELECT COUNT(*) FROM wallets WHERE chain = ?", (chain,)).fetchone()[0]
            return conn.execute("SELECT COUNT(*) FROM wallets").fetchone()[0]

    def record_rotation(self, wallet_id: str, new_address: str, reason: str = ""):
        with self._conn() as conn:
            conn.execute(
                "INSERT INTO rotations (wallet_id, new_address, rotated_at, reason) VALUES (?, ?, ?, ?)",
                (wallet_id, new_address, time.time(), reason),
            )
            conn.commit()

    def get_rotations(self, wallet_id: str) -> list[dict]:
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT * FROM rotations WHERE wallet_id = ? ORDER BY rotated_at DESC",
                (wallet_id,),
            ).fetchall()
            return [dict(r) for r in rows]

    def stats(self) -> dict[str, Any]:
        with self._conn() as conn:
            total = conn.execute("SELECT COUNT(*) FROM wallets").fetchone()[0]
            by_chain = conn.execute(
                "SELECT chain, COUNT(*) as count FROM wallets GROUP BY chain ORDER BY count DESC"
            ).fetchall()
            encrypted = conn.execute("SELECT COUNT(*) FROM wallets WHERE encrypted = 1").fetchone()[0]
            return {
                "total_wallets": total,
                "encrypted_wallets": encrypted,
                "chains_used": len(by_chain),
                "by_chain": {r["chain"]: r["count"] for r in by_chain},
                "rotations": conn.execute("SELECT COUNT(*) FROM rotations").fetchone()[0],
        }

    def _row_to_entry(self, row) -> WalletEntry:
        def _g(key: str, default=""):
            try:
                val = row[key]
                return val if val is not None else default
            except (KeyError, IndexError, TypeError):
                return default
        return WalletEntry(
            id=_g("id"),
            chain=_g("chain"),
            address=_g("address"),
            label=_g("label"),
            tags=json.loads(_g("tags", "[]")),
            group=_g("wallet_group"),
            created_at=float(_g("created_at", "0")),
            encrypted_key=_g("encrypted_key"),
            public_key=_g("public_key"),
            derivation_path=_g("derivation_path"),
            hd_path=_g("hd_path"),
            mnemonic_encrypted=_g("mnemonic_encrypted"),
            encrypted=bool(_g("encrypted", "0")),
            balance_usd=float(_g("balance_usd", "0")),
            notes=_g("notes"),
        )


_vault: Optional[Vault] = None


def get_vault() -> Vault:
    global _vault
    if _vault is None:
        _vault = Vault(cfg.db_path)
    return _vault
