import secrets
import string
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util.Padding import pad, unpad
from Crypto.Hash import SHA256
import base64
import gradio as gr

# Custom CSS for modern styling
custom_css = """
body {
    font-family: 'Inter', sans-serif;
    background: linear-gradient(135deg, #1e3c72 0%, #2a5298 100%);
}

.gradio-container {
    max-width: 800px;
    margin: 20px auto;
    background: rgba(255, 255, 255, 0.95);
    border-radius: 16px;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.2);
    padding: 24px;
}

h1, .gr-markdown h1 {
    color: #1e3c72;
    font-weight: 700;
    text-align: center;
    margin-bottom: 16px;
}

.gr-button {
    border-radius: 8px;
    padding: 12px 24px;
    font-weight: 600;
    transition: all 0.3s ease;
}

.gr-button-primary {
    background: #2a5298;
    color: white;
    border: none;
}

.gr-button-primary:hover {
    background: #1e3c72;
    transform: translateY(-2px);
}

.gr-button-secondary {
    background: #f1f3f5;
    color: #2a5298;
    border: 1px solid #2a5298;
}

.gr-button-secondary:hover {
    background: #e9ecef;
    transform: translateY(-2px);
}

.gr-textbox, .gr-radio {
    border-radius: 8px;
    border: 1px solid #dee2e6;
    background: #f8f9fa;
}

.gr-textbox:focus {
    border-color: #2a5298;
    box-shadow: 0 0 0 3px rgba(42, 82, 152, 0.2);
}

.gr-row {
    gap: 16px;
}

.card {
    background: white;
    border-radius: 12px;
    padding: 16px;
    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
    margin-bottom: 16px;
}

.footer {
    text-align: center;
    color: #6c757d;
    font-size: 0.9em;
    margin-top: 24px;
}
"""

def generate_random_key(length=32):
    """Generate a random key with digits, uppercase, lowercase, and special characters."""
    characters = string.ascii_lowercase + string.ascii_uppercase + string.digits + string.punctuation
    key = (
        secrets.choice(string.ascii_lowercase) +
        secrets.choice(string.ascii_uppercase) +
        secrets.choice(string.digits) +
        secrets.choice(string.punctuation)
    )
    key += ''.join(secrets.choice(characters) for _ in range(length - 4))
    key_list = list(key)
    secrets.SystemRandom().shuffle(key_list)
    return ''.join(key_list)

def derive_key(custom_key, key_size):
    """Derive a key of specified size (128, 192, or 256 bits) from a custom string using SHA-256."""
    hasher = SHA256.new()
    hasher.update(custom_key.encode('utf-8'))
    full_hash = hasher.digest()
    return full_hash[:int(key_size) // 8]

def encrypt_string(plain_text, key, key_size):
    """Encrypt a string using AES in CBC mode with specified key size."""
    try:
        if not plain_text:
            return "Error: Please provide text to encrypt."
        if not key:
            return "Error: Please provide a key."
        derived_key = derive_key(key, key_size)
        plain_text_bytes = plain_text.encode('utf-8')
        iv = get_random_bytes(AES.block_size)
        cipher = AES.new(derived_key, AES.MODE_CBC, iv)
        padded_text = pad(plain_text_bytes, AES.block_size)
        cipher_text = cipher.encrypt(padded_text)
        return base64.b64encode(iv + cipher_text).decode('utf-8')
    except Exception as e:
        return f"Encryption error: {str(e)}"

def decrypt_string(encrypted_text, key, key_size):
    """Decrypt a string using AES in CBC mode with specified key size."""
    try:
        if not encrypted_text:
            return "Error: Please provide encrypted text."
        if not key:
            return "Error: Please provide a key."
        derived_key = derive_key(key, key_size)
        raw = base64.b64decode(encrypted_text)
        iv = raw[:AES.block_size]
        cipher_text = raw[AES.block_size:]
        cipher = AES.new(derived_key, AES.MODE_CBC, iv)
        padded_text = cipher.decrypt(cipher_text)
        return unpad(padded_text, AES.block_size).decode('utf-8')
    except Exception as e:
        return f"Decryption error: {str(e)}"

def generate_and_display_key():
    """Generate and return a random key."""
    return generate_random_key()

# Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🔒 AES Encryption/Decryption App
        Securely encrypt or decrypt text using AES with 128, 192, or 256-bit keys.
        """
    )

    with gr.Column(elem_classes="card"):
        gr.Markdown("### Configuration")
        key_size = gr.Radio(
            choices=["128", "192", "256"],
            label="AES Key Size (bits)",
            value="256",
            elem_classes="radio"
        )
        
        with gr.Row():
            custom_key = gr.Textbox(
                label="Custom Key",
                placeholder="Enter your key or generate one",
                lines=1,
                show_copy_button=True
            )
            generate_key_btn = gr.Button("Generate Key", variant="secondary")

    with gr.Column(elem_classes="card"):
        gr.Markdown("### Input")
        input_text = gr.Textbox(
            label="Input Text",
            placeholder="Enter text to encrypt or encrypted text to decrypt",
            lines=4
        )
        
        with gr.Row():
            encrypt_btn = gr.Button("🔐 Encrypt", variant="primary")
            decrypt_btn = gr.Button("🔓 Decrypt", variant="secondary")

    with gr.Column(elem_classes="card"):
        gr.Markdown("### Output")
        output_text = gr.Textbox(
            label="Result",
            placeholder="Result will appear here",
            lines=4,
            interactive=False,
            show_copy_button=True
        )

    gr.Markdown(
        """
        <div class="footer">
            Powered by Gradio & PyCrypto | © 2025 SecureApp
        </div>
        """
    )

    # Event handlers
    generate_key_btn.click(
        fn=generate_and_display_key,
        outputs=custom_key
    )
    
    encrypt_btn.click(
        fn=encrypt_string,
        inputs=[input_text, custom_key, key_size],
        outputs=output_text
    )
    
    decrypt_btn.click(
        fn=decrypt_string,
        inputs=[input_text, custom_key, key_size],
        outputs=output_text
    )

# Launch the app
demo.launch()