from flask import Flask, render_template, request, jsonify
import hashlib
import hmac
import os
import re
import secrets
import shutil
import string
import sys
import tempfile
import threading
import time
from datetime import datetime, timedelta, timezone
from functools import wraps

import pysqlcipher3.dbapi2 as sqlcipher

app = Flask(__name__)
app.secret_key = secrets.token_hex(32)

DB_PATH = os.environ.get("PINVAULT_DB", "/data/pinvault.db")
LOCAL_BACKUP_DIR = os.environ.get("PINVAULT_LOCAL_BACKUP_DIR", "/data/backups")
NAS_BACKUP_DIR = os.environ.get("PINVAULT_NAS_BACKUP_DIR", "")
BACKUP_INTERVAL = int(os.environ.get("PINVAULT_BACKUP_INTERVAL", "3600"))
MAX_BACKUPS = int(os.environ.get("PINVAULT_MAX_BACKUPS", "168"))
RECOVERY_CODE_COUNT = 4
RECOVERY_CODE_LENGTH = 64
ALPHANUMERIC = string.ascii_letters + string.digits
API_KEY = os.environ.get("PINVAULT_API_KEY") or secrets.token_hex(32)

_pin_attempt_locks = {}
_PIN_MAX_ATTEMPTS = 5
_PIN_LOCKOUT_SECONDS = 900


def _check_pin_rate_limit(pin_id):
    now = time.monotonic()
    if pin_id not in _pin_attempt_locks:
        _pin_attempt_locks[pin_id] = {"attempts": 0, "lockout_until": 0}
    entry = _pin_attempt_locks[pin_id]
    if now < entry["lockout_until"]:
        remaining = int(entry["lockout_until"] - now)
        return False, remaining
    entry["attempts"] += 1
    if entry["attempts"] > _PIN_MAX_ATTEMPTS:
        entry["lockout_until"] = now + _PIN_LOCKOUT_SECONDS
        entry["attempts"] = 0
        return False, _PIN_LOCKOUT_SECONDS
    return True, 0


def _reset_pin_rate_limit(pin_id):
    if pin_id in _pin_attempt_locks:
        _pin_attempt_locks[pin_id] = {"attempts": 0, "lockout_until": 0}


def require_api_key(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        auth = request.headers.get("Authorization", "")
        if not auth.startswith("Bearer ") or auth[7:] != API_KEY:
            return jsonify({"error": "Unauthorized"}), 401
        return f(*args, **kwargs)
    return decorated


def dict_factory(cursor, row):
    d = {}
    for idx, col in enumerate(cursor.description):
        d[col[0].lower()] = row[idx]
    return d


def _derive_sqlcipher_key(bcrypt_hash: str) -> str:
    """Derive SQLCipher key from bcrypt hash using PBKDF2 (#5)."""
    salt = b"pinvault-sqlcipher-v1"  # Fixed salt - bcrypt hash provides entropy
    derived = hashlib.pbkdf2_hmac("sha256", bcrypt_hash.encode(), salt, 100_000)
    return derived.hex()[:64]


def get_conn():
    conn = sqlcipher.connect(DB_PATH)
    conn.row_factory = dict_factory
    conn.execute(f'PRAGMA key = "{app.master_key}"')
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("SELECT count(*) FROM sqlite_master")
    return conn


LEGACY_SALT = "pinvault-recovery-v1"


def hash_val(value, salt=None):
    """Hash a PIN or recovery code with HMAC-SHA256 using a per-row salt (#4)."""
    effective_salt = salt if salt else LEGACY_SALT
    return hmac.new(effective_salt.encode(), value.encode(), hashlib.sha256).hexdigest()


def generate_recovery_code():
    return "".join(secrets.choice(ALPHANUMERIC) for _ in range(RECOVERY_CODE_LENGTH))


def now_utc():
    return datetime.now(timezone.utc)


def init_db():
    with get_conn() as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS pins (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                label TEXT NOT NULL DEFAULT '',
                pin_hash TEXT NOT NULL,
                created_at TEXT NOT NULL,
                lock_until TEXT NOT NULL,
                bypassed_count INTEGER NOT NULL DEFAULT 0,
                revealed INTEGER NOT NULL DEFAULT 0
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS recovery_codes (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                pin_id INTEGER NOT NULL,
                code_hash TEXT NOT NULL,
                used INTEGER NOT NULL DEFAULT 0,
                FOREIGN KEY (pin_id) REFERENCES pins(id)
            )
        """)
        try:
            conn.execute("ALTER TABLE pins ADD COLUMN revealed INTEGER NOT NULL DEFAULT 0")
            conn.commit()
        except Exception:
            pass
        for table in ("pins", "recovery_codes"):
            try:
                conn.execute(f"ALTER TABLE {table} ADD COLUMN salt TEXT")
                conn.commit()
            except Exception:
                pass


def _backup_db(dest_path: str) -> bool:
    try:
        shutil.copy2(DB_PATH, dest_path)
        return True
    except Exception:
        if os.path.exists(dest_path):
            os.remove(dest_path)
        return False


def _atomic_copy(src, dst):
    """Copy src over dst atomically (temp file + rename) to avoid a torn DB."""
    dst_dir = os.path.dirname(dst) or "."
    fd, tmp = tempfile.mkstemp(dir=dst_dir, prefix=".pinvault-restore-")
    try:
        with os.fdopen(fd, "wb") as out, open(src, "rb") as f:
            shutil.copyfileobj(f, out)
        os.replace(tmp, dst)
    except BaseException:
        try:
            os.unlink(tmp)
        except OSError:
            pass
        raise


def _cleanup_backups(directory: str):
    try:
        files = sorted(
            [f for f in os.listdir(directory) if f.startswith("pinvault-backup-") and f.endswith(".db")],
            reverse=True,
        )
        for old in files[MAX_BACKUPS:]:
            os.remove(os.path.join(directory, old))
    except Exception:
        pass


def _backup_dirs():
    return [d for d in (LOCAL_BACKUP_DIR, NAS_BACKUP_DIR) if d]


def _do_backup():
    ts = now_utc().strftime("%Y%m%d-%H%M%S")
    filename = f"pinvault-backup-{ts}.db"

    os.makedirs(LOCAL_BACKUP_DIR, exist_ok=True)
    local_path = os.path.join(LOCAL_BACKUP_DIR, filename)

    ok = _backup_db(local_path)
    if ok:
        _cleanup_backups(LOCAL_BACKUP_DIR)
        if NAS_BACKUP_DIR:
            os.makedirs(NAS_BACKUP_DIR, exist_ok=True)
            nas_path = os.path.join(NAS_BACKUP_DIR, filename)
            shutil.copy2(local_path, nas_path)
            _cleanup_backups(NAS_BACKUP_DIR)
        return True
    return False


def _backup_loop():
    while True:
        try:
            _do_backup()
        except Exception as e:
            print(f"Backup failed: {e}", file=sys.stderr, flush=True)
        timer = threading.Timer(BACKUP_INTERVAL, _backup_loop)
        timer.daemon = True
        timer.start()
        break


def _get_backup_list():
    result = []
    for d in _backup_dirs():
        location = "local" if d == LOCAL_BACKUP_DIR else "nas"
        try:
            for f in sorted(os.listdir(d)):
                if f.startswith("pinvault-backup-") and f.endswith(".db"):
                    path = os.path.join(d, f)
                    size = os.path.getsize(path)
                    mtime = datetime.fromtimestamp(os.path.getmtime(path), tz=timezone.utc)
                    result.append({
                        "filename": f,
                        "location": location,
                        "size": size,
                        "created": mtime.isoformat(),
                    })
        except Exception:
            pass
    return sorted(result, key=lambda x: x["created"], reverse=True)


@app.route("/")
def index():
    return render_template("index.html")


@app.route("/api/status", methods=["GET"])
@require_api_key
def api_status():
    return jsonify({"locked": not getattr(app, "master_key", None)})


@app.route("/api/backups", methods=["GET"])
@require_api_key
def api_list_backups():
    if not getattr(app, "master_key", None):
        return jsonify({"error": "Vault not configured."}), 401
    return jsonify(_get_backup_list())


@app.route("/api/backups", methods=["POST"])
@require_api_key
def api_trigger_backup():
    if not getattr(app, "master_key", None):
        return jsonify({"error": "Vault not configured."}), 401
    ok = _do_backup()
    return jsonify({"ok": ok, "backups": _get_backup_list()})


@app.route("/api/backups/<filename>/restore", methods=["POST"])
@require_api_key
def api_restore_backup(filename):
    if not getattr(app, "master_key", None):
        return jsonify({"error": "Vault not configured."}), 401
    if not re.match(r'^pinvault-backup-\d{8}-\d{6}\.db$', filename):
        return jsonify({"error": "Invalid backup filename."}), 400

    src = None
    for d in _backup_dirs():
        candidate = os.path.join(d, filename)
        if os.path.exists(candidate):
            src = candidate
            break

    if not src:
        return jsonify({"error": "Backup file not found."}), 404

    try:
        test_conn = sqlcipher.connect(src)
        test_conn.execute(f'PRAGMA key = "{app.master_key}"')
        test_conn.execute("SELECT count(*) FROM sqlite_master")
        test_conn.close()
    except Exception as e:
        return jsonify({"error": f"Backup verification failed: {e}"}), 500

    _atomic_copy(src, DB_PATH)
    return jsonify({"ok": True, "restored": filename})


@app.route("/api/pins", methods=["GET"])
@require_api_key
def api_list_pins():
    if not getattr(app, "master_key", None):
        return jsonify({"error": "Vault not configured."}), 401
    with get_conn() as conn:
        rows = conn.execute(
            """
            SELECT p.id, p.label, p.created_at, p.lock_until, p.bypassed_count, p.revealed,
                   (SELECT COUNT(*) FROM recovery_codes rc WHERE rc.pin_id = p.id AND rc.used = 0) AS remaining_codes
            FROM pins p ORDER BY p.id
            """
        ).fetchall()
    now = now_utc()
    result = []
    for row in rows:
        lock_until = datetime.fromisoformat(row["lock_until"])
        locked = now < lock_until
        days_left = max(0, (lock_until - now).days) if locked else 0
        result.append({
            "id": row["id"],
            "label": row["label"],
            "created_at": row["created_at"],
            "lock_until": row["lock_until"],
            "locked": locked,
            "days_remaining": days_left,
            "revealed": bool(row["revealed"]),
            "bypassed_count": row["bypassed_count"],
            "remaining_codes": row["remaining_codes"],
        })
    return jsonify(result)


@app.route("/api/pins", methods=["POST"])
@require_api_key
def api_add_pin():
    if not getattr(app, "master_key", None):
        return jsonify({"error": "Vault not configured."}), 401
    data = request.get_json() or {}
    label = data.get("label", "").strip()
    lock_days = data.get("lock_days", 30)
    if isinstance(lock_days, bool) or not isinstance(lock_days, int) or lock_days < 0:
        return jsonify({"error": "lock_days must be an integer >= 0"}), 400

    pin = f"{secrets.randbelow(10000):04d}"
    pin_salt = secrets.token_hex(16)
    pin_hash = hash_val(pin, pin_salt)
    now = now_utc()
    lock_until = (now + timedelta(days=lock_days)).isoformat()

    codes = [generate_recovery_code() for _ in range(RECOVERY_CODE_COUNT)]

    with get_conn() as conn:
        cur = conn.execute(
            "INSERT INTO pins (label, pin_hash, salt, created_at, lock_until) VALUES (?, ?, ?, ?, ?)",
            (label, pin_hash, pin_salt, now.isoformat(), lock_until),
        )
        pin_id = cur.lastrowid
        for code in codes:
            code_salt = secrets.token_hex(16)
            conn.execute(
                "INSERT INTO recovery_codes (pin_id, code_hash, salt) VALUES (?, ?, ?)",
                (pin_id, hash_val(code, code_salt), code_salt),
            )
        conn.commit()

    return jsonify({
        "id": pin_id,
        "label": label,
        "pin": pin,
        "lock_until": lock_until,
        "recovery_codes": codes,
    }), 201


@app.route("/api/pins/<int:pin_id>/access", methods=["POST"])
@require_api_key
def api_access_pin(pin_id):
    if not getattr(app, "master_key", None):
        return jsonify({"error": "Vault not configured."}), 401
    allowed, remaining = _check_pin_rate_limit(pin_id)
    if not allowed:
        return jsonify({"error": f"Too many attempts. Try again in {remaining}s."}), 429
    data = request.get_json() or {}
    bypass_code = data.get("bypass_code", "").strip() if data.get("bypass_code") else None

    with get_conn() as conn:
        row = conn.execute("SELECT * FROM pins WHERE id = ?", (pin_id,)).fetchone()
        if not row:
            return jsonify({"error": "PIN entry not found"}), 404

        lock_until = datetime.fromisoformat(row["lock_until"])
        now = now_utc()

        if bypass_code:
            code_row = None
            for candidate in conn.execute(
                "SELECT * FROM recovery_codes WHERE pin_id = ? AND used = 0",
                (pin_id,),
            ).fetchall():
                if hmac.compare_digest(hash_val(bypass_code, candidate["salt"]), candidate["code_hash"]):
                    code_row = candidate
                    break
            if not code_row:
                return jsonify({"error": "Invalid or already-used recovery code"}), 403
            conn.execute(
                "UPDATE recovery_codes SET used = 1 WHERE id = ?",
                (code_row["id"],),
            )
            conn.execute(
                "UPDATE pins SET lock_until = ?, bypassed_count = bypassed_count + 1, revealed = 1 WHERE id = ?",
                (now.isoformat(), pin_id),
            )
            conn.commit()
            _reset_pin_rate_limit(pin_id)
            pin = recover_pin(row["pin_hash"], row["salt"])
            return jsonify({"pin": pin})

        if now >= lock_until:
            conn.execute("UPDATE pins SET revealed = 1 WHERE id = ?", (pin_id,))
            conn.commit()
            _reset_pin_rate_limit(pin_id)
            pin = recover_pin(row["pin_hash"], row["salt"])
            return jsonify({"pin": pin})

        return jsonify({"error": "PIN is locked. Use a recovery code to bypass."}), 423

    return jsonify({"error": "Unexpected error"}), 500


@app.route("/api/pins/<int:pin_id>", methods=["DELETE"])
@require_api_key
def api_delete_pin(pin_id):
    if not getattr(app, "master_key", None):
        return jsonify({"error": "Vault not configured."}), 401
    with get_conn() as conn:
        row = conn.execute("SELECT revealed FROM pins WHERE id = ?", (pin_id,)).fetchone()
        if not row:
            return jsonify({"error": "PIN entry not found"}), 404
        if not row["revealed"]:
            return jsonify({"error": "Cannot delete unrevealed PIN. Access it first."}), 409
        conn.execute("DELETE FROM recovery_codes WHERE pin_id = ?", (pin_id,))
        conn.execute("DELETE FROM pins WHERE id = ?", (pin_id,))
        conn.commit()
    return jsonify({"ok": True})


def recover_pin(pin_hash, salt=None):
    """Recover a 4-digit PIN by exhaustive search (design: PINs are 10^4 space)."""
    for i in range(10000):
        candidate = f"{i:04d}"
        if hmac.compare_digest(hash_val(candidate, salt), pin_hash):
            return candidate
    return None


def _verify_db():
    try:
        conn = sqlcipher.connect(DB_PATH)
        conn.execute(f'PRAGMA key = "{app.master_key}"')
        conn.execute("SELECT count(*) FROM sqlite_master")
        conn.close()
        return True
    except Exception:
        return False


def _latest_backup():
    candidates = []
    for d in _backup_dirs():
        try:
            for f in os.listdir(d):
                if f.startswith("pinvault-backup-") and f.endswith(".db"):
                    path = os.path.join(d, f)
                    candidates.append((os.path.getmtime(path), path))
        except Exception:
            pass
    if not candidates:
        return None
    candidates.sort(reverse=True)
    return candidates[0][1]


def _auto_restore():
    if _verify_db():
        return True

    print("DB is corrupted or unreadable. Attempting auto-restore...", flush=True)
    backup = _latest_backup()
    if not backup:
        print("No backups found. Cannot auto-restore.", file=sys.stderr, flush=True)
        return False

    try:
        test = sqlcipher.connect(backup)
        test.execute(f'PRAGMA key = "{app.master_key}"')
        test.execute("SELECT count(*) FROM sqlite_master")
        test.close()
    except Exception:
        print(f"Backup {backup} also unreadable. Cannot auto-restore.", file=sys.stderr, flush=True)
        return False

    _atomic_copy(backup, DB_PATH)
    print(f"Restored DB from backup: {backup}", flush=True)
    return True


_bootstrap_ok = False


def bootstrap():
    global _bootstrap_ok
    bcrypt_hash = os.environ.get("PINVAULT_MASTER_HASH", "").strip()
    os.environ.pop("PINVAULT_MASTER_HASH", None)

    if not bcrypt_hash:
        print("ERROR: PINVAULT_MASTER_HASH environment variable is required.", file=sys.stderr)
        print("Generate: python3 -c \"import bcrypt; print(bcrypt.hashpw(b'YOUR_PASSWORD', bcrypt.gensalt()).decode())\"", file=sys.stderr)
        return False

    if not bcrypt_hash.startswith("$2"):
        print("ERROR: PINVAULT_MASTER_HASH must be a bcrypt hash (starts with $2b$ or $2a$)", file=sys.stderr)
        return False

    app.master_key = _derive_sqlcipher_key(bcrypt_hash)

    if os.path.exists(DB_PATH):
        if not _auto_restore():
            print("FATAL: Cannot open database and no valid backup found.", file=sys.stderr, flush=True)
            return False

    init_db()
    _do_backup()
    t = threading.Thread(target=_backup_loop, daemon=True)
    t.start()
    _bootstrap_ok = True
    return True


_bootstrap_ok = bootstrap()

if __name__ == "__main__":
    if not _bootstrap_ok:
        sys.exit(1)
    app.run(host="0.0.0.0", port=8765)