from datetime import datetime, timedelta
from jose import jwt
from passlib.context import CryptContext

SECRET_KEY = "your-secret-key"
ALGORITHM = "HS256"

pwd_context = CryptContext(
    schemes=["argon2", "bcrypt"],  # 添加 argon2 作为备选方案
    default="argon2",  # 使用 argon2 作为默认方案
    deprecated="auto",
)

def hash_password(password: str) -> str:
    return pwd_context.hash(password)


def verify_password(plain_password, hashed_password):
    return pwd_context.verify(plain_password, hashed_password)

def create_access_token(data: dict):
    to_encode = data.copy()
    expire = datetime.utcnow() + timedelta(minutes=30)
    to_encode.update({"exp": expire})
    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)

from Crypto.PublicKey import RSA

def generate_key_pair():
    key = RSA.generate(2048)
    private_key = key.export_key().decode('utf-8')
    public_key = key.publickey().export_key().decode('utf-8')
    return private_key, public_key

from Crypto.Cipher import PKCS1_OAEP, AES
from Crypto.Hash import SHA256
from Crypto.Random import get_random_bytes
import base64

def encrypt_data(data: str, public_key_pem: str) -> str:
    # 1. 准备 RSA 公钥
    recipient_key = RSA.import_key(public_key_pem)
    cipher_rsa = PKCS1_OAEP.new(recipient_key)

    # 2. 生成随机 AES 会话密钥 (32 bytes)
    session_key = get_random_bytes(32)

    # 3. 用 RSA 加密 AES 密钥
    enc_session_key = cipher_rsa.encrypt(session_key)

    # 4. 用 AES-GCM 加密实际数据
    cipher_aes = AES.new(session_key, AES.MODE_GCM)
    ciphertext, tag = cipher_aes.encrypt_and_digest(data.encode('utf-8'))

    # 5. 组合结果: [RSA加密的密钥(256B)] + [Nonce(16B)] + [Tag(16B)] + [密文]
    combined = enc_session_key + cipher_aes.nonce + tag + ciphertext
    return base64.b64encode(combined).decode('utf-8')

def decrypt_data(encrypted_data_b64: str, private_key_pem: str) -> str:
    # 1. 解码合并的二进制数据
    encrypted_data = base64.b64decode(encrypted_data_b64)
    
    # 2. 分解组件 (RSA 2048位产生的密文固定为 256 字节)
    rsa_key_len = 256 
    enc_session_key = encrypted_data[:rsa_key_len]
    nonce = encrypted_data[rsa_key_len : rsa_key_len + 16]
    tag = encrypted_data[rsa_key_len + 16 : rsa_key_len + 32]
    ciphertext = encrypted_data[rsa_key_len + 32 :]

    # 3. 用 RSA 私钥解密 AES 会话密钥
    private_key = RSA.import_key(private_key_pem)
    cipher_rsa = PKCS1_OAEP.new(private_key)
    session_key = cipher_rsa.decrypt(enc_session_key)

    # 4. 用 AES 会话密钥解密数据
    cipher_aes = AES.new(session_key, AES.MODE_GCM, nonce=nonce)
    dec_data = cipher_aes.decrypt_and_verify(ciphertext, tag)
    
    return dec_data.decode('utf-8')

from fastapi import Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer
from sqlalchemy.future import select
from sqlalchemy.ext.asyncio import AsyncSession
from jose import JWTError
from . import database, models

oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

async def get_current_user(token: str = Depends(oauth2_scheme), db: AsyncSession = Depends(database.get_db)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
    except JWTError:
        raise credentials_exception
    
    result = await db.execute(select(models.User).where(models.User.username == username))
    user = result.scalars().first()
    if user is None:
        raise credentials_exception
    return user