from flask import Flask, render_template, request, jsonify
import hashlib
import hmac
import os
import random
import secrets
import shutil
import string
import sys
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 not hmac.compare_digest(auth[7:], API_KEY):
            return jsonify({"error": "Invalid or missing API key"}), 401
        return f(*args, **kwargs)
    return decorated
    d = {}
    for idx, col in enumerate(cursor.description):
        d[col[0]] = row[idx]
    return d


def _derive_sqlcipher_key(bcrypt_hash: str) -> str:
    return hashlib.sha256(bcrypt_hash.encode()).hexdigest()[: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


def hash_val(value):
    return hashlib.sha256(value.encode()).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


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 _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 _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)
    os.makedirs(NAS_BACKUP_DIR, exist_ok=True)

    local_path = os.path.join(LOCAL_BACKUP_DIR, filename)
    nas_path = os.path.join(NAS_BACKUP_DIR, filename)

    ok = _backup_db(local_path)
    if ok:
        shutil.copy2(local_path, nas_path)
        _cleanup_backups(LOCAL_BACKUP_DIR)
        _cleanup_backups(NAS_BACKUP_DIR)
        return ok
    return False


def _backup_loop():
    while True:
        _do_backup()
        timer = threading.Timer(BACKUP_INTERVAL, _backup_loop)
        timer.daemon = True
        timer.start()
        break


def _get_backup_list():
    result = []
    for d in [LOCAL_BACKUP_DIR, NAS_BACKUP_DIR]:
        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
    import re
    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 [LOCAL_BACKUP_DIR, NAS_BACKUP_DIR]:
        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

    shutil.copy2(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()
    label = data.get("label", "").strip()
    lock_days = data.get("lock_days", 30)
    if lock_days < 0:
        return jsonify({"error": "lock_days must be >= 0"}), 400

    pin = f"{random.randint(0, 9999):04d}"
    pin_hash = hash_val(pin)
    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, created_at, lock_until) VALUES (?, ?, ?, ?)",
            (label, pin_hash, now.isoformat(), lock_until),
        )
        pin_id = cur.lastrowid
        for code in codes:
            conn.execute(
                "INSERT INTO recovery_codes (pin_id, code_hash) VALUES (?, ?)",
                (pin_id, hash_val(code)),
            )
        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_hash = hash_val(bypass_code)
            code_row = conn.execute(
                "SELECT * FROM recovery_codes WHERE pin_id = ? AND code_hash = ? AND used = 0",
                (pin_id, code_hash),
            ).fetchone()
            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"])
            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"])
            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,))
        cur = conn.execute("DELETE FROM pins WHERE id = ?", (pin_id,))
        conn.commit()
    return jsonify({"ok": True})


def recover_pin(pin_hash):
    for i in range(10000):
        candidate = f"{i:04d}"
        if hash_val(candidate) == 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 [LOCAL_BACKUP_DIR, NAS_BACKUP_DIR]:
        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

    shutil.copy2(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)