import re
import time

from flask import Blueprint, request, jsonify, g
from app import db, limiter
from app.models.user import User
from app.models.audit_log import AuditLog
from app.services.auth_service import (
    hash_auth_token,
    verify_auth_token,
    generate_tokens,
    generate_mfa_token,
    decode_token,
    blacklist_token,
    require_jwt,
    encrypt_totp_secret,
    decrypt_totp_secret,
)

auth_bp = Blueprint('auth', __name__)

EMAIL_RE = re.compile(r'^[^@\s]+@[^@\s]+\.[^@\s]+$')


def _client_ip():
    """Return the best-effort client IP from the request context."""
    return request.headers.get('X-Forwarded-For', request.remote_addr or '').split(',')[0].strip()


@auth_bp.route('/register', methods=['POST'])
@limiter.limit('10 per minute')
def register():
    data = request.get_json(silent=True) or {}
    email = (data.get('email') or '').strip().lower()
    auth_hash = data.get('auth_hash', '')
    enc_key_salt = data.get('enc_key_salt', '')

    if not email or not EMAIL_RE.match(email):
        return jsonify({'error': 'Invalid email address'}), 400
    if not auth_hash:
        return jsonify({'error': 'auth_hash is required'}), 400
    if not enc_key_salt:
        return jsonify({'error': 'enc_key_salt is required'}), 400

    if User.query.filter_by(email=email).first():
        return jsonify({'error': 'Email already registered'}), 409

    master_hash = hash_auth_token(auth_hash)
    user = User(email=email, master_hash=master_hash, enc_key_salt=enc_key_salt)
    db.session.add(user)
    db.session.flush()  # populate user.id before logging

    AuditLog.log(
        user_id=user.id,
        action='auth.register',
        resource_type='user',
        resource_id=user.id,
        detail=f'New account registered: {email}',
        ip_address=_client_ip(),
    )
    db.session.commit()

    return jsonify({'message': 'Account created successfully'}), 201


@auth_bp.route('/login', methods=['POST'])
@limiter.limit('10 per minute')
def login():
    data = request.get_json(silent=True) or {}
    email = (data.get('email') or '').strip().lower()
    auth_hash = data.get('auth_hash', '')

    time.sleep(0.1)  # mitigate timing-based user enumeration

    if not email or not auth_hash:
        return jsonify({'error': 'Email and auth_hash are required'}), 400

    user = User.query.filter_by(email=email).first()
    if not user or not verify_auth_token(auth_hash, user.master_hash):
        if user:
            AuditLog.log(
                user_id=user.id,
                action='auth.login_failed',
                resource_type='user',
                resource_id=user.id,
                detail='Failed login attempt — invalid password',
                ip_address=_client_ip(),
            )
            db.session.commit()
        return jsonify({'error': 'Invalid email or password'}), 401

    from datetime import datetime
    user.last_login = datetime.utcnow()

    AuditLog.log(
        user_id=user.id,
        action='auth.login',
        resource_type='user',
        resource_id=user.id,
        detail=f'Successful login{" (MFA pending)" if user.totp_enabled else ""}',
        ip_address=_client_ip(),
    )
    db.session.commit()

    # MFA gate: if enabled, issue a short-lived mfa_token instead of full tokens
    if user.totp_enabled:
        mfa_token = generate_mfa_token(user.id)
        return jsonify({
            'mfa_required': True,
            'mfa_token': mfa_token,
            'enc_key_salt': user.enc_key_salt,
        }), 200

    tokens = generate_tokens(user.id)
    return jsonify({
        'access_token': tokens['access_token'],
        'refresh_token': tokens['refresh_token'],
        'enc_key_salt': user.enc_key_salt,
    }), 200


@auth_bp.route('/logout', methods=['POST'])
def logout():
    """Blacklist both the access token (from header) and refresh token (from body)."""
    auth_header = request.headers.get('Authorization', '')
    if auth_header.startswith('Bearer '):
        blacklist_token(auth_header[7:], 'access')

    data = request.get_json(silent=True) or {}
    refresh_token = data.get('refresh_token', '')
    if refresh_token:
        blacklist_token(refresh_token, 'refresh')

    return jsonify({'message': 'Logged out'}), 200


@auth_bp.route('/refresh', methods=['POST'])
@limiter.limit('30 per minute')
def refresh():
    data = request.get_json(silent=True) or {}
    refresh_token = data.get('refresh_token', '')
    if not refresh_token:
        return jsonify({'error': 'refresh_token is required'}), 400

    try:
        payload = decode_token(refresh_token, expected_type='refresh')
    except Exception:
        return jsonify({'error': 'Invalid or expired refresh token'}), 401

    # Rotate: blacklist old refresh token and issue fresh pair
    blacklist_token(refresh_token, 'refresh')
    tokens = generate_tokens(int(payload['sub']))
    return jsonify({
        'access_token': tokens['access_token'],
        'refresh_token': tokens['refresh_token'],
    }), 200


# ── MFA / TOTP endpoints ─────────────────────────────────────────────────────

@auth_bp.route('/mfa/setup', methods=['GET'])
@require_jwt
def mfa_setup():
    """Generate a new TOTP secret and return QR code (as base64 PNG data URI)."""
    user = db.session.get(User, g.current_user_id)
    if user.totp_enabled:
        return jsonify({'error': 'MFA is already enabled'}), 400

    import pyotp
    import qrcode
    import io
    import base64

    secret = pyotp.random_base32()
    uri = pyotp.TOTP(secret).provisioning_uri(
        name=user.email,
        issuer_name='PassKeeper',
    )

    img = qrcode.make(uri)
    buf = io.BytesIO()
    img.save(buf, format='PNG')
    qr_b64 = base64.b64encode(buf.getvalue()).decode()

    return jsonify({
        'secret': secret,
        'qr_code': f'data:image/png;base64,{qr_b64}',
        'uri': uri,
    }), 200


@auth_bp.route('/mfa/enable', methods=['POST'])
@require_jwt
def mfa_enable():
    """Enable MFA after verifying the first TOTP code."""
    user = db.session.get(User, g.current_user_id)
    if user.totp_enabled:
        return jsonify({'error': 'MFA is already enabled'}), 400

    data = request.get_json(silent=True) or {}
    secret = (data.get('secret') or '').strip()
    totp_code = (data.get('totp_code') or '').strip()

    if not secret or not totp_code:
        return jsonify({'error': 'secret and totp_code are required'}), 400

    import pyotp
    if not pyotp.TOTP(secret).verify(totp_code, valid_window=1):
        return jsonify({'error': 'Invalid verification code'}), 400

    totp_secret_enc, totp_iv = encrypt_totp_secret(secret)
    user.totp_secret = totp_secret_enc
    user.totp_iv = totp_iv
    user.totp_enabled = True

    AuditLog.log(
        user_id=user.id,
        action='auth.mfa_enable',
        resource_type='user',
        resource_id=user.id,
        detail='TOTP two-factor authentication enabled',
        ip_address=_client_ip(),
    )
    db.session.commit()

    return jsonify({'message': 'MFA enabled successfully'}), 200


@auth_bp.route('/mfa/disable', methods=['POST'])
@require_jwt
def mfa_disable():
    """Disable MFA after verifying the current TOTP code."""
    user = db.session.get(User, g.current_user_id)
    if not user.totp_enabled:
        return jsonify({'error': 'MFA is not enabled'}), 400

    data = request.get_json(silent=True) or {}
    totp_code = (data.get('totp_code') or '').strip()

    import pyotp
    plaintext_secret = decrypt_totp_secret(user.totp_secret, user.totp_iv)
    if not pyotp.TOTP(plaintext_secret).verify(totp_code, valid_window=1):
        return jsonify({'error': 'Invalid verification code'}), 400

    user.totp_secret = None
    user.totp_iv = None
    user.totp_enabled = False

    AuditLog.log(
        user_id=user.id,
        action='auth.mfa_disable',
        resource_type='user',
        resource_id=user.id,
        detail='TOTP two-factor authentication disabled',
        ip_address=_client_ip(),
    )
    db.session.commit()

    return jsonify({'message': 'MFA disabled'}), 200


@auth_bp.route('/mfa/verify', methods=['POST'])
@limiter.limit('10 per minute')
def mfa_verify():
    """Complete MFA login: verify TOTP code and exchange mfa_token for real tokens."""
    data = request.get_json(silent=True) or {}
    mfa_token = data.get('mfa_token', '')
    totp_code = (data.get('totp_code') or '').strip()

    if not mfa_token or not totp_code:
        return jsonify({'error': 'mfa_token and totp_code are required'}), 400

    try:
        payload = decode_token(mfa_token, expected_type='mfa', check_blacklist=True)
    except Exception:
        return jsonify({'error': 'Invalid or expired MFA token'}), 401

    user = db.session.get(User, int(payload['sub']))
    if not user or not user.totp_enabled:
        return jsonify({'error': 'MFA not configured for this account'}), 400

    import pyotp
    plaintext_secret = decrypt_totp_secret(user.totp_secret, user.totp_iv)
    if not pyotp.TOTP(plaintext_secret).verify(totp_code, valid_window=1):
        return jsonify({'error': 'Invalid verification code'}), 400

    # One-time use: blacklist the mfa_token
    blacklist_token(mfa_token, 'mfa')

    AuditLog.log(
        user_id=user.id,
        action='auth.mfa_verify',
        resource_type='user',
        resource_id=user.id,
        detail='MFA verification successful — session tokens issued',
        ip_address=_client_ip(),
    )
    db.session.commit()

    tokens = generate_tokens(user.id)
    return jsonify({
        'access_token': tokens['access_token'],
        'refresh_token': tokens['refresh_token'],
    }), 200


@auth_bp.route('/mfa/status', methods=['GET'])
@require_jwt
def mfa_status():
    user = db.session.get(User, g.current_user_id)
    return jsonify({'totp_enabled': user.totp_enabled}), 200


@auth_bp.route('/me', methods=['GET'])
@require_jwt
def me():
    """Return basic profile info for the authenticated user."""
    user = db.session.get(User, g.current_user_id)
    return jsonify({
        'id': user.id,
        'email': user.email,
        'created_at': user.created_at.isoformat() if user.created_at else None,
        'last_login': user.last_login.isoformat() if user.last_login else None,
        'totp_enabled': user.totp_enabled,
        'recovery_configured': bool(user.recovery_enc_salt),
    }), 200


# ── Account management ────────────────────────────────────────────────────────

@auth_bp.route('/change-password', methods=['POST'])
@require_jwt
def change_password():
    """
    Change master password — zero-knowledge atomic re-encryption.

    The client must:
      1. Derive current auth_hash and verify it locally against what it knows.
      2. Re-encrypt every vault item with the new vault key client-side.
      3. POST the new credentials + all re-encrypted item blobs in one request.

    The server verifies the current password, updates master_hash + enc_key_salt,
    and bulk-replaces all vault item ciphertexts atomically. If any step fails,
    the entire transaction is rolled back — the vault is never left in a split state.
    """
    data = request.get_json(silent=True) or {}
    current_auth_hash = data.get('current_auth_hash', '')
    new_auth_hash = data.get('new_auth_hash', '')
    new_enc_key_salt = data.get('new_enc_key_salt', '')
    items = data.get('items', [])  # [{id, enc_data, iv}, ...]

    if not current_auth_hash or not new_auth_hash or not new_enc_key_salt:
        return jsonify({'error': 'current_auth_hash, new_auth_hash, and new_enc_key_salt are required'}), 400

    user = db.session.get(User, g.current_user_id)

    if not verify_auth_token(current_auth_hash, user.master_hash):
        AuditLog.log(
            user_id=user.id,
            action='auth.change_password_failed',
            resource_type='user',
            resource_id=user.id,
            detail='Password change rejected — current password incorrect',
            ip_address=_client_ip(),
        )
        db.session.commit()
        return jsonify({'error': 'Current password is incorrect'}), 401

    try:
        from app.models.vault_item import VaultItem

        # Bulk-update all vault item ciphertexts with new vault key encryption
        item_ids = [i.get('id') for i in items if i.get('id')]
        existing = {
            v.id: v
            for v in VaultItem.query.filter(
                VaultItem.user_id == user.id,
                VaultItem.id.in_(item_ids),
            ).all()
        } if item_ids else {}

        for item_data in items:
            item_id = item_data.get('id')
            enc_data = item_data.get('enc_data', '')
            iv = item_data.get('iv', '')
            if not item_id or not enc_data or not iv:
                continue
            vault_item = existing.get(item_id)
            if vault_item:
                vault_item.enc_data = enc_data
                vault_item.iv = iv

        # Update credentials
        user.master_hash = hash_auth_token(new_auth_hash)
        user.enc_key_salt = new_enc_key_salt
        # Clear recovery data — it was encrypted with the old vault key and is now invalid
        user.recovery_enc_salt = None
        user.recovery_iv = None

        AuditLog.log(
            user_id=user.id,
            action='auth.change_password',
            resource_type='user',
            resource_id=user.id,
            detail=f'Master password changed; {len(existing)} vault item(s) re-encrypted; recovery code cleared',
            ip_address=_client_ip(),
        )
        db.session.commit()
    except Exception as e:
        db.session.rollback()
        return jsonify({'error': f'Password change failed: {str(e)}'}), 500

    return jsonify({'message': 'Password changed successfully. Please log in again.'}), 200


@auth_bp.route('/account', methods=['DELETE'])
@require_jwt
def delete_account():
    """
    Permanently delete the authenticated user's account and all associated data.
    Requires the current auth_hash for confirmation.
    Cascading deletes handle vault_items, folders, shared_items, emergency_access.
    """
    data = request.get_json(silent=True) or {}
    auth_hash = data.get('auth_hash', '')

    if not auth_hash:
        return jsonify({'error': 'auth_hash is required for account deletion'}), 400

    user = db.session.get(User, g.current_user_id)

    if not verify_auth_token(auth_hash, user.master_hash):
        AuditLog.log(
            user_id=user.id,
            action='auth.delete_account_failed',
            resource_type='user',
            resource_id=user.id,
            detail='Account deletion rejected — password incorrect',
            ip_address=_client_ip(),
        )
        db.session.commit()
        return jsonify({'error': 'Incorrect password'}), 401

    user_id = user.id
    user_email = user.email
    try:
        # Log before delete (user row will be gone after commit)
        AuditLog.log(
            user_id=user_id,
            action='auth.delete_account',
            resource_type='user',
            resource_id=user_id,
            detail=f'Account permanently deleted: {user_email}',
            ip_address=_client_ip(),
        )
        db.session.delete(user)
        db.session.commit()
    except Exception as e:
        db.session.rollback()
        return jsonify({'error': f'Account deletion failed: {str(e)}'}), 500

    return jsonify({'message': 'Account deleted'}), 200


# ── Account Recovery ──────────────────────────────────────────────────────────

@auth_bp.route('/recovery/setup', methods=['POST'])
@require_jwt
def recovery_setup():
    """
    Store a recovery-key-encrypted copy of enc_key_salt.

    The client generates a random 128-bit recovery code, derives a recovery key
    from it (PBKDF2), encrypts enc_key_salt with that key (AES-256-GCM), and
    sends the ciphertext + iv. The server stores these blobs — it never sees the
    recovery code or enc_key_salt plaintext.

    The recovery code is displayed to the user once and never stored server-side.
    """
    data = request.get_json(silent=True) or {}
    recovery_enc_salt = data.get('recovery_enc_salt', '').strip()
    recovery_iv = data.get('recovery_iv', '').strip()

    if not recovery_enc_salt or not recovery_iv:
        return jsonify({'error': 'recovery_enc_salt and recovery_iv are required'}), 400

    user = db.session.get(User, g.current_user_id)
    user.recovery_enc_salt = recovery_enc_salt
    user.recovery_iv = recovery_iv

    AuditLog.log(
        user_id=user.id,
        action='auth.recovery_setup',
        resource_type='user',
        resource_id=user.id,
        detail='Account recovery code configured',
        ip_address=_client_ip(),
    )
    db.session.commit()

    return jsonify({'message': 'Recovery code saved'}), 200


@auth_bp.route('/recovery/status', methods=['GET'])
@require_jwt
def recovery_status():
    """Return whether the user has a recovery code configured."""
    user = db.session.get(User, g.current_user_id)
    return jsonify({'recovery_configured': bool(user.recovery_enc_salt)}), 200


@auth_bp.route('/recover', methods=['POST'])
@limiter.limit('5 per minute')
def recover_account():
    """
    Recover account access using a recovery code.

    Flow:
      1. Client looks up enc_key_salt and recovery blobs by email.
      2. Client decrypts enc_key_salt using the recovery key (derived from the recovery code).
      3. Client derives new auth_hash and new vault key with a new master password.
      4. Client re-encrypts all vault items with the new vault key.
      5. Client POSTs everything here in one atomic payload.

    This endpoint is unauthenticated — the recovery code is the credential.
    """
    data = request.get_json(silent=True) or {}
    email = (data.get('email') or '').strip().lower()
    new_auth_hash = data.get('new_auth_hash', '')
    new_enc_key_salt = data.get('new_enc_key_salt', '')
    # Proof that the client successfully decrypted enc_key_salt:
    # the client re-derives auth_hash from the *original* enc_key_salt path
    # and sends it alongside the new credentials for server-side verification.
    recovery_proof = data.get('recovery_proof', '')
    items = data.get('items', [])

    if not all([email, new_auth_hash, new_enc_key_salt, recovery_proof]):
        return jsonify({'error': 'email, new_auth_hash, new_enc_key_salt, and recovery_proof are required'}), 400

    time.sleep(0.1)  # timing mitigation

    user = User.query.filter_by(email=email).first()
    if not user or not user.recovery_enc_salt:
        return jsonify({'error': 'No recovery code found for this account'}), 404

    # recovery_proof is the enc_key_salt re-encrypted by the client using the
    # recovery key — we return it as a blob for the client to verify, then
    # the client sends back the decrypted enc_key_salt as recovery_proof.
    # Simpler: recovery_proof = HMAC or simply the decrypted enc_key_salt itself,
    # which the client proves by sending it back plaintext. The server checks it
    # matches user.enc_key_salt — if the recovery code was wrong, decryption
    # would produce garbage that won't match.
    if recovery_proof != user.enc_key_salt:
        AuditLog.log(
            user_id=user.id,
            action='auth.recovery_failed',
            resource_type='user',
            resource_id=user.id,
            detail='Recovery attempt failed — incorrect recovery code',
            ip_address=_client_ip(),
        )
        db.session.commit()
        return jsonify({'error': 'Invalid recovery code'}), 401

    try:
        from app.models.vault_item import VaultItem

        item_ids = [i.get('id') for i in items if i.get('id')]
        existing = {
            v.id: v
            for v in VaultItem.query.filter(
                VaultItem.user_id == user.id,
                VaultItem.id.in_(item_ids),
            ).all()
        } if item_ids else {}

        for item_data in items:
            item_id = item_data.get('id')
            enc_data = item_data.get('enc_data', '')
            iv = item_data.get('iv', '')
            if not item_id or not enc_data or not iv:
                continue
            vault_item = existing.get(item_id)
            if vault_item:
                vault_item.enc_data = enc_data
                vault_item.iv = iv

        user.master_hash = hash_auth_token(new_auth_hash)
        user.enc_key_salt = new_enc_key_salt
        # Recovery code is consumed — clear it so it cannot be reused
        user.recovery_enc_salt = None
        user.recovery_iv = None

        AuditLog.log(
            user_id=user.id,
            action='auth.recovery_success',
            resource_type='user',
            resource_id=user.id,
            detail=f'Account recovered; {len(existing)} vault item(s) re-encrypted; recovery code consumed',
            ip_address=_client_ip(),
        )
        db.session.commit()
    except Exception as e:
        db.session.rollback()
        return jsonify({'error': f'Recovery failed: {str(e)}'}), 500

    tokens = generate_tokens(user.id)
    return jsonify({
        'message': 'Account recovered successfully',
        'access_token': tokens['access_token'],
        'refresh_token': tokens['refresh_token'],
        'enc_key_salt': user.enc_key_salt,
    }), 200


@auth_bp.route('/recovery/data', methods=['GET'])
@limiter.limit('10 per minute')
def recovery_data():
    """
    Return the data the client needs to attempt recovery (unauthenticated).
    Exposes only: enc_key_salt, recovery_enc_salt, recovery_iv.
    Returns 404 if no recovery code is configured (prevents user enumeration
    of which accounts have recovery set up — same response for unknown email).
    """
    email = (request.args.get('email') or '').strip().lower()
    if not email:
        return jsonify({'error': 'email is required'}), 400

    user = User.query.filter_by(email=email).first()
    if not user or not user.recovery_enc_salt:
        return jsonify({'error': 'No recovery data found'}), 404

    return jsonify({
        'enc_key_salt': user.enc_key_salt,
        'recovery_enc_salt': user.recovery_enc_salt,
        'recovery_iv': user.recovery_iv,
    }), 200


@auth_bp.route('/recovery/items', methods=['GET'])
@limiter.limit('10 per minute')
def recovery_items():
    """
    Return encrypted vault items for recovery re-encryption (unauthenticated).

    Requires X-Recovery-Proof header containing the plaintext enc_key_salt.
    The server verifies it matches user.enc_key_salt — proof that the client
    correctly decrypted the recovery blob (i.e. has the correct recovery code).

    Items are returned as encrypted ciphertext blobs only — no sensitive
    plaintext is exposed. The client re-encrypts them locally.
    """
    email = (request.args.get('email') or '').strip().lower()
    recovery_proof = request.headers.get('X-Recovery-Proof', '').strip()

    if not email or not recovery_proof:
        return jsonify({'error': 'email and X-Recovery-Proof header are required'}), 400

    user = User.query.filter_by(email=email).first()
    if not user or not user.recovery_enc_salt:
        return jsonify({'error': 'No recovery data found'}), 404

    if recovery_proof != user.enc_key_salt:
        AuditLog.log(
            user_id=user.id,
            action='auth.recovery_items_denied',
            resource_type='user',
            resource_id=user.id,
            detail='Recovery items request denied — incorrect recovery proof',
            ip_address=_client_ip(),
        )
        db.session.commit()
        return jsonify({'error': 'Invalid recovery proof'}), 401

    from app.models.vault_item import VaultItem
    items = VaultItem.query.filter_by(user_id=user.id).all()
    return jsonify({
        'items': [
            {'id': item.id, 'enc_data': item.enc_data, 'iv': item.iv}
            for item in items
        ]
    }), 200