from flask import Flask, render_template, request, jsonify
import hashlib
import os
import random
import secrets
import shutil
import string
import threading
from datetime import datetime, timedelta, timezone

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 = "/data/backups"
NAS_BACKUP_DIR = "/mnt/aidata/pinvault"
BACKUP_INTERVAL = 3600
MAX_BACKUPS = 168
RECOVERY_CODE_COUNT = 4
RECOVERY_CODE_LENGTH = 64
ALPHANUMERIC = string.ascii_letters + string.digits


def dict_factory(cursor, row):
    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"])
def api_status():
    return jsonify({"locked": not getattr(app, "master_key", None)})


@app.route("/api/backups", methods=["GET"])
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"])
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"])
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"])
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"])
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"])
def api_access_pin(pin_id):
    if not getattr(app, "master_key", None):
        return jsonify({"error": "Vault not configured."}), 401
    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()
            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()
            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"])
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


def bootstrap():
    import sys
    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 one with: python3 -c \"import bcrypt; print(bcrypt.hashpw(b'YOUR_PASSWORD', bcrypt.gensalt()).decode())\"", file=sys.stderr)
        sys.exit(1)

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

    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)
            sys.exit(1)

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


bootstrap()

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8765)