import gradio as gr
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse
from pathlib import Path
import hashlib
import os
import sqlite3
import tempfile
import zipfile
from datetime import datetime, timedelta
import base64
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.fernet import Fernet, InvalidToken
from apscheduler.schedulers.background import BackgroundScheduler

# ===================== CONFIG =====================
DATA_DIR = Path("/data")
ENCRYPTED_DIR = DATA_DIR / "encrypted"
ENCRYPTED_DIR.mkdir(parents=True, exist_ok=True)
DB_PATH = DATA_DIR / "files.db"

MAX_STORAGE_TB = 2
DEFAULT_TTL = 30
DEFAULT_DOWNLOAD_LIMIT = 10

# ===================== DB =====================
def init_db():
    conn = sqlite3.connect(DB_PATH)
    conn.execute('''
        CREATE TABLE IF NOT EXISTS files (
            sha TEXT PRIMARY KEY,
            original_name TEXT,
            upload_time TEXT,
            ttl_days INTEGER,
            download_limit INTEGER,
            downloads INTEGER DEFAULT 0,
            is_public INTEGER DEFAULT 0,
            size INTEGER
        )
    ''')
    conn.commit()
    conn.close()

init_db()

def get_db():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn

# ===================== CRYPTO =====================
def derive_key(password: str, salt: bytes = None):
    if salt is None:
        salt = os.urandom(16)
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=600_000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
    return key, salt

def encrypt_file(input_path: Path, password: str | None, is_public: bool):
    data = input_path.read_bytes()
    salt = os.urandom(16)
    
    if is_public:
        key, salt = derive_key("public-default-2026", salt)
    else:
        key, salt = derive_key(password, salt)
    
    fernet = Fernet(key)
    encrypted = fernet.encrypt(data)
    
    sha = hashlib.sha256(encrypted).hexdigest()
    enc_path = ENCRYPTED_DIR / f"{sha}.enc"
    enc_path.write_bytes(salt + encrypted)
    
    return enc_path, encrypted

def decrypt_file(sha: str, password: str | None) -> bytes:
    enc_path = ENCRYPTED_DIR / f"{sha}.enc"
    if not enc_path.exists():
        raise FileNotFoundError("File not found")
    
    data = enc_path.read_bytes()
    salt = data[:16]
    encrypted = data[16:]
    
    if password is None:  # Public
        key, _ = derive_key("public-default-2026", salt)
    else:
        key, _ = derive_key(password, salt)
    
    fernet = Fernet(key)
    return fernet.decrypt(encrypted)

# ===================== HELPERS =====================
def get_usage():
    total_size = sum(f.stat().st_size for f in ENCRYPTED_DIR.glob("*.enc") if f.is_file())
    file_count = len(list(ENCRYPTED_DIR.glob("*.enc")))
    return total_size, file_count

def cleanup_old_files():
    now = datetime.now()
    conn = get_db()
    rows = conn.execute("SELECT sha, upload_time, ttl_days FROM files").fetchall()
    for row in rows:
        expiry = datetime.fromisoformat(row['upload_time']) + timedelta(days=row['ttl_days'])
        if now > expiry:
            (ENCRYPTED_DIR / f"{row['sha']}.enc").unlink(missing_ok=True)
            conn.execute("DELETE FROM files WHERE sha=?", (row['sha'],))
    conn.commit()
    conn.close()

scheduler = BackgroundScheduler()
scheduler.add_job(cleanup_old_files, 'interval', hours=4)
scheduler.start()

# ===================== UPLOAD =====================
def handle_upload(files, password: str, confirm_pw: str, ttl_days: int, download_limit: int, is_public: bool, progress=gr.Progress()):
    if not files:
        return "❌ No files selected.", None, None
    
    if not is_public and (not password or password != confirm_pw):
        return "❌ Private files require matching password and confirmation.", None, None
    
    progress(0.1, desc="Preparing...")
    
    if isinstance(files, list) and len(files) > 1:
        with tempfile.TemporaryDirectory() as tmpdir:
            zip_path = Path(tmpdir) / "archive.zip"
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
                for i, f in enumerate(files):
                    progress(0.3 + 0.6 * i / len(files), desc=f"Zipping {Path(f.name).name}")
                    zf.write(f.name, Path(f.name).name)
            orig_name = "archive.zip"
            enc_path, _ = encrypt_file(zip_path, password if not is_public else None, is_public)
    else:
        f = files[0] if isinstance(files, list) else files
        orig_name = Path(f.name).name
        progress(0.4, desc="Encrypting...")
        enc_path, _ = encrypt_file(Path(f.name), password if not is_public else None, is_public)
    
    sha = enc_path.stem
    size = enc_path.stat().st_size
    
    conn = get_db()
    conn.execute('''
        INSERT OR REPLACE INTO files 
        (sha, original_name, upload_time, ttl_days, download_limit, downloads, is_public, size)
        VALUES (?, ?, ?, ?, ?, 0, ?, ?)
    ''', (sha, orig_name, datetime.now().isoformat(), ttl_days, download_limit, int(is_public), size))
    conn.commit()
    conn.close()
    
    used, count = get_usage()
    status = f"✅ Uploaded!\nSHA: {sha[:16]}...\nStorage: {used/(1024**4):.3f} TB / {MAX_STORAGE_TB} TB"
    space_id = os.getenv('SPACE_ID', 'your-space')
    link = f"https://{space_id}.hf.space/dw/{sha}"
    return status, sha, link

# ===================== DOWNLOAD =====================
def download_file(sha: str, deck: str = None):
    conn = get_db()
    row = conn.execute("SELECT * FROM files WHERE sha=?", (sha,)).fetchone()
    if not row:
        raise HTTPException(status_code=404, detail="File not found")
    
    if row['downloads'] >= row['download_limit'] and row['download_limit'] > 0:
        (ENCRYPTED_DIR / f"{sha}.enc").unlink(missing_ok=True)
        conn.execute("DELETE FROM files WHERE sha=?", (sha,))
        conn.commit()
        raise HTTPException(status_code=410, detail="File expired (download limit reached)")
    
    try:
        decrypted = decrypt_file(sha, None if row['is_public'] else deck)
        conn.execute("UPDATE files SET downloads = downloads + 1 WHERE sha=?", (sha,))
        conn.commit()
    except InvalidToken:
        raise HTTPException(status_code=401, detail="Invalid password")
    except Exception:
        raise HTTPException(status_code=500, detail="Decryption failed")
    finally:
        conn.close()
    
    orig_name = row['original_name']
    media_type = "application/zip" if orig_name.endswith(".zip") else "application/octet-stream"
    
    return StreamingResponse(
        iter([decrypted]),
        media_type=media_type,
        headers={"Content-Disposition": f'attachment; filename="{orig_name}"'}
    )

# ===================== GRADIO UI =====================
with gr.Blocks(title="Secure Vault") as demo:
    gr.Markdown("# 🔒 Secure Encrypted File Vault\n**2 TB Storage • Auto-expiry • Passwords never stored**")
    
    with gr.Tab("📤 Upload"):
        with gr.Row():
            files_in = gr.File(label="Select one or multiple files", file_count="multiple")
            
            with gr.Column():
                is_public_in = gr.Checkbox(label="🌍 Public (no password)", value=False)
                pw_in = gr.Textbox(label="🔑 Encryption Password", type="password")
                confirm_pw_in = gr.Textbox(label="🔑 Confirm Password", type="password")
                ttl_in = gr.Slider(1, 30, value=DEFAULT_TTL, step=1, label="⏳ TTL (days)")
                dlimit_in = gr.Slider(1, 50, value=DEFAULT_DOWNLOAD_LIMIT, step=1, label="🔢 Max Downloads")
        
        upload_btn = gr.Button("🔐 Encrypt & Upload", variant="primary")
        
        upload_status = gr.Textbox(label="Status", lines=3)
        share_link_out = gr.Textbox(label="🔗 Shareable Download Link")
        
        def toggle_visibility(is_public):
            return gr.update(visible=not is_public), gr.update(visible=not is_public)
        
        is_public_in.change(toggle_visibility, is_public_in, [pw_in, confirm_pw_in])
        
        upload_btn.click(
            handle_upload,
            inputs=[files_in, pw_in, confirm_pw_in, ttl_in, dlimit_in, is_public_in],
            outputs=[upload_status, gr.State(), share_link_out]
        )
    
    with gr.Tab("📁 File Browser"):
        with gr.Row():
            search_in = gr.Textbox(placeholder="Search...", label="Filter")
            refresh_btn = gr.Button("Refresh")
        
        file_table = gr.Dataframe(
            headers=["SHA", "Name", "Uploaded", "Expires In", "Downloads", "Public"],
            value=[],
            interactive=False
        )
        
        def refresh_browser(search_term=""):
            conn = get_db()
            rows = conn.execute("SELECT * FROM files ORDER BY upload_time DESC").fetchall()
            conn.close()
            
            data = []
            for r in rows:
                expiry_days = max(0, (datetime.fromisoformat(r['upload_time']) + timedelta(days=r['ttl_days']) - datetime.now()).days)
                data.append([
                    r['sha'][:16] + "...",
                    r['original_name'],
                    r['upload_time'][:10],
                    f"{expiry_days} days",
                    f"{r['downloads']}/{r['download_limit']}",
                    "✅ Yes" if r['is_public'] else "🔒 No"
                ])
            return data
        
        refresh_btn.click(refresh_browser, inputs=search_in, outputs=file_table)
    
    with gr.Tab("📊 Dashboard"):
        usage_md = gr.Markdown()
        
        def update_usage():
            used, count = get_usage()
            pct = (used / (MAX_STORAGE_TB * 1024**4)) * 100
            return f"""
**Storage**: {used / (1024**4):.3f} TB / {MAX_STORAGE_TB} TB ({pct:.1f}%)  
**Files**: {count}  
**Status**: {'🟢 OK' if pct < 85 else '⚠️ Near limit'}
            """
        
        demo.load(update_usage, outputs=usage_md)
    
    gr.Markdown("""
    ### Download Instructions
    • Public: `https://your-space.hf.space/dw/{SHA}`  
    • Private: `https://your-space.hf.space/dw/{SHA}?deck=YOUR_PASSWORD`
    """)

# ===================== FASTAPI + MOUNT =====================
fastapi_app = FastAPI(title="Secure Vault")

@fastapi_app.get("/dw/{sha}")
async def dw_route(sha: str, request: Request):
    deck = request.query_params.get("deck")
    return download_file(sha, deck)

# Mount Gradio
app = gr.mount_gradio_app(fastapi_app, demo, path="/")

# For HF Spaces - do NOT run uvicorn here