import os
import json
import base64
from app import app, db
from models.user import User
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

def int_to_b64url(val):
    byte_len = (val.bit_length() + 7) // 8
    b = val.to_bytes(byte_len, 'big')
    return base64.urlsafe_b64encode(b).decode('utf-8').rstrip('=')

def seed_admin_user(email="benjimanmiller@gmail.com", password="0kl4h0m4", first_name="Benjiman", last_name="Miller"):
    with app.app_context():
        # 1. Generate RSA-OAEP 2048 keypair
        priv = rsa.generate_private_key(public_exponent=65537, key_size=2048)
        pub = priv.public_key()
        priv_nums = priv.private_numbers()
        pub_nums = pub.public_numbers()

        pub_jwk = {
            'kty': 'RSA',
            'n': int_to_b64url(pub_nums.n),
            'e': int_to_b64url(pub_nums.e),
            'alg': 'RSA-OAEP-256',
            'ext': True,
            'key_ops': ['encrypt']
        }

        priv_jwk = {
            'kty': 'RSA',
            'n': int_to_b64url(pub_nums.n),
            'e': int_to_b64url(pub_nums.e),
            'd': int_to_b64url(priv_nums.d),
            'p': int_to_b64url(priv_nums.p),
            'q': int_to_b64url(priv_nums.q),
            'dp': int_to_b64url(priv_nums.dmp1),
            'dq': int_to_b64url(priv_nums.dmq1),
            'qi': int_to_b64url(priv_nums.iqmp),
            'alg': 'RSA-OAEP-256',
            'ext': True,
            'key_ops': ['decrypt']
        }

        # 2. Derive KEK via PBKDF2
        salt = os.urandom(16)
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000
        )
        kek = kdf.derive(password.encode('utf-8'))

        # 3. Encrypt private key via AES-GCM
        iv = os.urandom(12)
        aesgcm = AESGCM(kek)
        ct = aesgcm.encrypt(iv, json.dumps(priv_jwk).encode('utf-8'), None)

        enc_bundle = json.dumps({
            'iv': base64.b64encode(iv).decode('utf-8'),
            'ct': base64.b64encode(ct).decode('utf-8')
        })
        salt_b64 = base64.b64encode(salt).decode('utf-8')
        pub_str = json.dumps(pub_jwk)

        # 4. Check if user exists
        user = User.query.filter_by(email=email).first()
        if user:
            print(f"Updating existing user '{email}' to admin with new keys and password...")
            user.first_name = first_name
            user.last_name = last_name
            user.is_admin = True
            user.is_banned = False
            user.set_password(password)
            user.public_key = pub_str
            user.encrypted_private_key = enc_bundle
            user.encryption_salt = salt_b64
        else:
            print(f"Creating new admin user '{email}'...")
            user = User(
                email=email,
                first_name=first_name,
                last_name=last_name,
                is_admin=True,
                is_banned=False,
                public_key=pub_str,
                encrypted_private_key=enc_bundle,
                encryption_salt=salt_b64
            )
            user.set_password(password)
            db.session.add(user)

        db.session.commit()
        print(f"Successfully seeded admin user: {email} (is_admin={user.is_admin})")

if __name__ == '__main__':
    seed_admin_user()