# backend/vault/credential_vault.py
import os
import json
import sqlite3
import time
import uuid
import logging

try:
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
except ImportError:
    AESGCM = None

try:
    from argon2.low_level import hash_secret_raw, Type
except ImportError:
    hash_secret_raw = None
    Type = None

class CredentialVault:
    def __init__(self, db_path: str):
        # db_path = "{app_data_dir}/vault.db" — encrypted at rest
        os.makedirs(os.path.dirname(os.path.abspath(db_path)), exist_ok=True)
        self.db = sqlite3.connect(db_path, check_same_thread=False)
        self.db.execute('PRAGMA journal_mode=WAL')
        self._init_schema()
        # Schema: credentials(id, service_name, account_identifier,
        #         encrypted_token_blob, nonce, salt, created_at, updated_at)

    def _init_schema(self):
        cursor = self.db.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS credentials (
                id TEXT PRIMARY KEY,
                service_name TEXT NOT NULL,
                account_identifier TEXT NOT NULL,
                encrypted_token_blob BLOB NOT NULL,
                nonce BLOB NOT NULL,
                salt BLOB NOT NULL,
                created_at REAL NOT NULL,
                updated_at REAL NOT NULL,
                UNIQUE(service_name, account_identifier)
            )
        ''')
        self.db.commit()

    def _derive_key(self, master_password: str, salt: bytes) -> bytes:
        # Argon2id — real KDF, not MD5/SHA1
        if hash_secret_raw is None:
            raise RuntimeError("argon2-cffi is not installed")
            
        return hash_secret_raw(
            secret=master_password.encode(),
            salt=salt,
            time_cost=3,
            memory_cost=65536,
            parallelism=4,
            hash_len=32,
            type=Type.ID
        )

    def store(self, service: str, account: str, token_data: dict, master_password: str):
        if AESGCM is None:
            raise RuntimeError("cryptography is not installed")
            
        salt = os.urandom(16)
        key = self._derive_key(master_password, salt)
        nonce = os.urandom(12)
        aesgcm = AESGCM(key)
        
        plaintext = json.dumps(token_data).encode('utf-8')
        ciphertext = aesgcm.encrypt(nonce, plaintext, None)
        
        # Store salt+ciphertext in db — key never stored, only re-derived on access
        cursor = self.db.cursor()
        now = time.time()
        
        cursor.execute('''
            INSERT INTO credentials (id, service_name, account_identifier, encrypted_token_blob, nonce, salt, created_at, updated_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(service_name, account_identifier) DO UPDATE SET
                encrypted_token_blob = excluded.encrypted_token_blob,
                nonce = excluded.nonce,
                salt = excluded.salt,
                updated_at = excluded.updated_at
        ''', (str(uuid.uuid4()), service, account, ciphertext, nonce, salt, now, now))
        
        self.db.commit()
        logging.info(f"Stored encrypted credentials for {service}/{account}")

    def retrieve(self, service: str, account: str, master_password: str) -> dict:
        if AESGCM is None:
            raise RuntimeError("cryptography is not installed")
            
        cursor = self.db.cursor()
        cursor.execute('''
            SELECT encrypted_token_blob, nonce, salt
            FROM credentials
            WHERE service_name = ? AND account_identifier = ?
        ''', (service, account))
        
        row = cursor.fetchone()
        if not row:
            raise KeyError(f"No credentials found for {service}/{account}")
            
        ciphertext, nonce, salt = row
        key = self._derive_key(master_password, salt)
        aesgcm = AESGCM(key)
        
        try:
            plaintext = aesgcm.decrypt(nonce, ciphertext, None)
            return json.loads(plaintext.decode('utf-8'))
        except Exception as e:
            logging.error(f"Failed to decrypt credentials for {service}/{account}. Wrong password or corrupted data.")
            raise ValueError("Decryption failed") from e

    def list_accounts(self) -> list[tuple[str, str]]:
        cursor = self.db.cursor()
        cursor.execute('SELECT service_name, account_identifier FROM credentials')
        return cursor.fetchall()

    def delete(self, service: str, account: str):
        cursor = self.db.cursor()
        cursor.execute('''
            DELETE FROM credentials 
            WHERE service_name = ? AND account_identifier = ?
        ''', (service, account))
        self.db.commit()