from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from app.db.session import get_db
from app.core.deps import get_current_active_user
from app.models.user import User
from app.services.encryption_service import encryption_service
from pydantic import BaseModel

router = APIRouter()

class KeyPairResponse(BaseModel):
    public_key: str
    private_key: str

class EncryptRequest(BaseModel):
    message: str
    recipient_user_id: int

class EncryptResponse(BaseModel):
    encrypted_message: str

class DecryptRequest(BaseModel):
    encrypted_message: str
    private_key: str

class DecryptResponse(BaseModel):
    decrypted_message: str

@router.post("/generate-keys", response_model=KeyPairResponse)
def generate_encryption_keys(
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db)
):
    """Generate new RSA key pair for E2E encryption"""
    private_key, public_key = encryption_service.generate_rsa_key_pair()
    
    # Update user's public key in database
    current_user.public_key = public_key
    db.commit()
    
    return KeyPairResponse(
        public_key=public_key,
        private_key=private_key
    )

@router.get("/public-key")
def get_public_key(
    current_user: User = Depends(get_current_active_user)
):
    """Get current user's public key"""
    return {
        "public_key": current_user.public_key,
        "user_id": current_user.id
    }

@router.get("/public-key/{user_id}")
def get_user_public_key(
    user_id: int,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db)
):
    """Get another user's public key"""
    user = db.query(User).filter(User.id == user_id).first()
    if not user:
        raise HTTPException(status_code=404, detail="User not found")
    
    return {
        "public_key": user.public_key,
        "user_id": user.id,
        "username": user.username
    }

@router.post("/encrypt", response_model=EncryptResponse)
def encrypt_message(
    request: EncryptRequest,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db)
):
    """Encrypt a message for a specific user"""
    recipient = db.query(User).filter(User.id == request.recipient_user_id).first()
    if not recipient:
        raise HTTPException(status_code=404, detail="Recipient not found")
    
    if not recipient.public_key:
        raise HTTPException(status_code=400, detail="Recipient has no public key")
    
    encrypted_message = encryption_service.encrypt_message(
        request.message, 
        recipient.public_key
    )
    
    return EncryptResponse(encrypted_message=encrypted_message)

@router.post("/decrypt", response_model=DecryptResponse)
def decrypt_message(
    request: DecryptRequest,
    current_user: User = Depends(get_current_active_user)
):
    """Decrypt a message using private key"""
    try:
        decrypted_message = encryption_service.decrypt_message(
            request.encrypted_message,
            request.private_key
        )
        return DecryptResponse(decrypted_message=decrypted_message)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Decryption failed: {str(e)}")

@router.put("/public-key")
def update_public_key(
    public_key: str,
    current_user: User = Depends(get_current_active_user),
    db: Session = Depends(get_db)
):
    """Update user's public key"""
    current_user.public_key = public_key
    db.commit()
    
    return {"message": "Public key updated successfully"}