import gradio as gr
import sqlite3
import hashlib
import os
import base64
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

# --- Database Setup ---
DB_PATH = "vault_system.db"

def init_db():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    # Users table stores hashed master passwords (we never store plain text)
    cursor.execute('''CREATE TABLE IF NOT EXISTS users 
                      (username TEXT PRIMARY KEY, master_hash TEXT, salt BLOB)''')
    # Vault table stores encrypted secrets
    cursor.execute('''CREATE TABLE IF NOT EXISTS vault 
                      (username TEXT, service TEXT, encrypted_secret TEXT)''')
    conn.commit()
    conn.close()

init_db()

# --- Security Helper Functions ---

def derive_key(password, salt):
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
    return Fernet(key)

def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()

# --- App Functions ---

def handle_auth(username, password, mode):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    if mode == "Register":
        salt = os.urandom(16)
        m_hash = hash_password(password)
        try:
            cursor.execute("INSERT INTO users VALUES (?, ?, ?)", (username, m_hash, salt))
            conn.commit()
            return "✅ Account Created! Please Login."
        except:
            return "❌ Username already exists."
            
    else: # Login
        cursor.execute("SELECT master_hash, salt FROM users WHERE username=?", (username,))
        result = cursor.fetchone()
        if result and result[0] == hash_password(password):
            return "SUCCESS"
        return "❌ Invalid Username or Password."

def vault_operation(action, user, master_pwd, service=None, secret=None):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # Get user salt to derive the same key
    cursor.execute("SELECT salt FROM users WHERE username=?", (user,))
    salt = cursor.fetchone()[0]
    cipher = derive_key(master_pwd, salt)

    if action == "Save":
        enc_secret = cipher.encrypt(secret.encode()).decode()
        cursor.execute("INSERT INTO vault VALUES (?, ?, ?)", (user, service, enc_secret))
        conn.commit()
        return f"✅ Saved {service} encrypted."
    
    elif action == "Get":
        cursor.execute("SELECT encrypted_secret FROM vault WHERE username=? AND service=?", (user, service))
        res = cursor.fetchone()
        if res:
            decrypted = cipher.decrypt(res[0].encode()).decode()
            return f"🔑 {service} Password: {decrypted}"
        return "❌ Service not found."

# --- UI with Custom CSS/JS ---
custom_style = """
.container { max-width: 800px; margin: auto; }
.login-card { border: 1px solid #e0e0e0; padding: 40px; border-radius: 15px; background: #fff; box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
.nav-header { text-align: center; border-bottom: 2px solid #8A2BE2; margin-bottom: 20px; }
"""

with gr.Blocks(css=custom_style, title="SecureVault Pro") as demo:
    user_session = gr.State("")
    pwd_session = gr.State("")

    # --- LOGIN PAGE ---
    with gr.Column(visible=True) as login_page:
        gr.HTML("<div class='nav-header'><h1>🔐 SecureVault Manager</h1><p>Professional AES-256 Encryption Suite</p></div>")
        with gr.Row(elem_classes="login-card"):
            with gr.Column():
                u_in = gr.Textbox(label="Username")
                p_in = gr.Textbox(label="Master Password", type="password")
                auth_mode = gr.Radio(["Login", "Register"], value="Login", label="Action")
                auth_btn = gr.Button("Access Vault", variant="primary")
                auth_msg = gr.Markdown()

    # --- DASHBOARD PAGE ---
    with gr.Column(visible=False) as dash_page:
        gr.HTML("<div class='nav-header'><h1>💼 My Secure Vault</h1></div>")
        
        with gr.Tabs():
            with gr.Tab("➕ Add Secret"):
                svc = gr.Textbox(label="Service Name (e.g. GitHub)")
                val = gr.Textbox(label="Password", type="password")
                save_btn = gr.Button("Encrypt & Store", variant="primary")
            
            with gr.Tab("🔍 Retrieve"):
                search = gr.Textbox(label="Service Name")
                get_btn = gr.Button("Decrypt & Reveal", variant="secondary")
        
        op_msg = gr.Textbox(label="System Output", interactive=False)
        logout_btn = gr.Button("🔒 Secure Logout", size="sm")

    # --- Event Logic ---
    def on_auth(u, p, m):
        status = handle_auth(u, p, m)
        if status == "SUCCESS":
            return gr.update(visible=False), gr.update(visible=True), u, p, ""
        return gr.update(visible=True), gr.update(visible=False), "", "", status

    auth_btn.click(on_auth, [u_in, p_in, auth_mode], [login_page, dash_page, user_session, pwd_session, auth_msg])
    
    save_btn.click(lambda u, p, s, v: vault_operation("Save", u, p, s, v), 
                   [user_session, pwd_session, svc, val], op_msg)
    
    get_btn.click(lambda u, p, s: vault_operation("Get", u, p, s), 
                  [user_session, pwd_session, search], op_msg)
    
    logout_btn.click(lambda: [gr.update(visible=True), gr.update(visible=False), "", ""], 
                     None, [login_page, dash_page, user_session, pwd_session])

if __name__ == "__main__":
    demo.launch()