import os
import json
import base64
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

# --- Helper: Generate AES Key ---
def generate_key(password: str, salt: bytes) -> bytes:
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    return kdf.derive(password.encode())

# --- Encrypt password with AES-GCM ---
def encrypt_password(password: str, key: bytes):
    # AES-GCM uses a 12-byte Nonce instead of a 16-byte IV
    nonce = os.urandom(12)
    aesgcm = AESGCM(key)
    # The encrypt method returns: ciphertext + 16-byte authentication tag
    ciphertext = aesgcm.encrypt(nonce, password.encode(), None)
    # We store Nonce + Ciphertext (which includes the Tag)
    return base64.b64encode(nonce + ciphertext).decode()

# --- Decrypt password with AES-GCM ---
def decrypt_password(encrypted_password: str, key: bytes):
    raw = base64.b64decode(encrypted_password)
    nonce, ciphertext = raw[:12], raw[12:]
    aesgcm = AESGCM(key)
    # Decrypt will automatically verify the Tag. If tampered with, it throws an InvalidTag exception.
    return aesgcm.decrypt(nonce, ciphertext, None).decode()

# --- Store password ---
def store_password(account: str, username: str, password: str, master_password: str):
    salt = os.urandom(16)
    key = generate_key(master_password, salt)
    encrypted_pw = encrypt_password(password, key)
    data = {"account": account, "username": username, "password": encrypted_pw, "salt": base64.b64encode(salt).decode()}
    
    try:
        with open("passwords.json", "r") as f:
            records = json.load(f)
    except FileNotFoundError:
        records = []

    records = [r for r in records if r["account"] != account]
    records.append(data)

    with open("passwords.json", "w") as f:
        json.dump(records, f, indent=4)
    print(f"[+] Authenticated GCM password stored for: {account}")

# --- Retrieve password ---
def retrieve_password(account: str, master_password: str):
    try:
        with open("passwords.json", "r") as f:
            records = json.load(f)
    except FileNotFoundError:
        return None

    for r in records:
        if r["account"] == account:
            salt = base64.b64decode(r["salt"])
            key = generate_key(master_password, salt)
            try:
                return decrypt_password(r["password"], key)
            except Exception:
                print("[-] Master password incorrect or data tampered with!")
                return None
    return None

# --- Encrypt file ---
def encrypt_file(file_path: str, account: str, master_password: str):
    password = retrieve_password(account, master_password)
    if not password: return

    try:
        with open(file_path, "rb") as f:
            content = f.read()
    except FileNotFoundError: return

    salt = os.urandom(16)
    nonce = os.urandom(12)
    key = generate_key(password, salt)
    
    aesgcm = AESGCM(key)
    ciphertext = aesgcm.encrypt(nonce, content, None)

    enc_file = file_path + ".enc"
    with open(enc_file, "wb") as f:
        f.write(ciphertext)

    meta = {
        "salt": base64.b64encode(salt).decode(),
        "nonce": base64.b64encode(nonce).decode(), # Now storing Nonce instead of IV
        "account": account
    }
    with open(enc_file + ".meta", "w") as f:
        json.dump(meta, f)

    print(f"[+] File encrypted with AES-GCM: {enc_file}")

# --- Decrypt file ---
def decrypt_file(enc_file: str, master_password: str):
    meta_file = enc_file + ".meta"
    if not os.path.exists(meta_file): return

    with open(meta_file, "r") as f:
        meta = json.load(f)

    account = meta["account"]
    salt = base64.b64decode(meta["salt"])
    nonce = base64.b64decode(meta["nonce"])

    password = retrieve_password(account, master_password)
    if not password: return

    key = generate_key(password, salt)
    with open(enc_file, "rb") as f:
        ciphertext = f.read()

    aesgcm = AESGCM(key)
    try:
        data = aesgcm.decrypt(nonce, ciphertext, None)
        dec_file = enc_file.replace(".enc", "_decrypted")
        with open(dec_file, "wb") as f:
            f.write(data)
        print(f"[+] Authenticated decryption successful: {dec_file}")
    except Exception:
        print("[-] Decryption failed: Invalid password or corrupted data.")