#!/usr/bin/env python3
"""
Decrypt WeChat SQLCipher databases to plaintext SQLite files.

Requirements:
    brew install sqlcipher

Usage:
    wechat-mcp-decrypt
    wechat-mcp-decrypt -o ./decrypted
"""

import json
import os
import subprocess
import sys
import glob
import argparse

from wechat_mcp.config import get_keys_path, get_db_base

PAGE_SZ = 4096
SALT_SZ = 16


def find_db_dir():
    db_base = get_db_base()
    pattern = os.path.join(db_base, "*", "db_storage")
    candidates = glob.glob(pattern)
    if len(candidates) >= 1:
        return candidates[0]
    if os.path.isdir(db_base) and os.path.basename(db_base) == "db_storage":
        return db_base
    return None


def find_sqlcipher():
    brew_path = "/opt/homebrew/opt/sqlcipher/bin/sqlcipher"
    if os.path.isfile(brew_path):
        return brew_path
    for p in os.environ.get("PATH", "").split(os.pathsep):
        candidate = os.path.join(p, "sqlcipher")
        if os.path.isfile(candidate):
            return candidate
    return None


def decrypt_database(sqlcipher_bin, src_path, dst_path, key_hex):
    """Decrypt a SQLCipher database to a plaintext SQLite file."""
    os.makedirs(os.path.dirname(dst_path), exist_ok=True)

    if os.path.exists(dst_path):
        os.remove(dst_path)

    sql_commands = f"""PRAGMA key = "x'{key_hex}'";
PRAGMA cipher_page_size = 4096;
ATTACH DATABASE '{dst_path}' AS plaintext KEY '';
SELECT sqlcipher_export('plaintext');
DETACH DATABASE plaintext;
"""

    try:
        result = subprocess.run(
            [sqlcipher_bin, src_path],
            input=sql_commands,
            capture_output=True,
            text=True,
            timeout=120,
        )
        if result.returncode != 0 or "Error" in result.stderr:
            return False, result.stderr.strip()

        if not os.path.isfile(dst_path) or os.path.getsize(dst_path) == 0:
            return False, "output file is empty"

        return True, "OK"
    except subprocess.TimeoutExpired:
        return False, "timeout"
    except Exception as e:
        return False, str(e)


def main():
    parser = argparse.ArgumentParser(description="Decrypt WeChat databases")
    parser.add_argument(
        "--keys", default=None,
        help="Path to wechat_keys.json (auto-detected if not specified)",
    )
    parser.add_argument(
        "-o", "--output", default="decrypted",
        help="Output directory for decrypted databases (default: decrypted)",
    )
    args = parser.parse_args()

    keys_file = args.keys or get_keys_path()
    if not os.path.isfile(keys_file):
        print(f"[-] Key file not found: {keys_file}")
        sys.exit(1)

    with open(keys_file, "r") as f:
        data = json.load(f)

    sqlcipher_bin = find_sqlcipher()
    if not sqlcipher_bin:
        print("[-] sqlcipher not found. Install it with: brew install sqlcipher")
        sys.exit(1)
    print(f"[*] Using sqlcipher: {sqlcipher_bin}")

    db_dir = find_db_dir()
    if not db_dir:
        print(f"[-] Could not find db_storage directory")
        sys.exit(1)
    print(f"[*] DB storage: {db_dir}")

    entries = {k: v for k, v in data.items() if not k.startswith("__")}
    print(f"[*] Decrypting {len(entries)} databases to {args.output}/\n")

    passed = 0
    failed = 0

    for db_rel_path, key_hex in sorted(entries.items()):
        src = os.path.join(db_dir, db_rel_path)
        dst = os.path.join(args.output, db_rel_path)

        if not os.path.isfile(src):
            print(f"  skip {db_rel_path}: source file not found")
            continue

        success, detail = decrypt_database(sqlcipher_bin, src, dst, key_hex)
        if success:
            dst_size = os.path.getsize(dst)
            print(f"  ok   {db_rel_path} -> {dst} ({dst_size / 1024:.0f} KB)")
            passed += 1
        else:
            print(f"  FAIL {db_rel_path}: {detail}")
            failed += 1

    print(f"\n[*] Done: {passed} decrypted, {failed} failed")
    if passed > 0:
        print(f"[*] Decrypted files saved to: {os.path.abspath(args.output)}/")


if __name__ == "__main__":
    main()
