import os
import sqlite3
import shutil
import mimetypes
from datetime import datetime, timezone
import gradio as gr

DB_NAME = "files.db"
UPLOAD_DIR = "uploads"
MAX_EMBEDDED_SIZE = 1 * 1024 * 1024  # 1 MB in bytes

def init_db(db_path=DB_NAME):
    """Initialize the SQLite database schema if it doesn't exist."""
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS files (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            filename TEXT NOT NULL,
            file_type TEXT,
            file_size_bytes INTEGER NOT NULL,
            upload_timestamp TEXT NOT NULL,
            is_large_file INTEGER NOT NULL,
            file_content_blob BLOB,
            relative_path TEXT
        )
    """)
    conn.commit()
    conn.close()
    os.makedirs(UPLOAD_DIR, exist_ok=True)

def process_uploaded_files(files, db_path=DB_NAME, upload_dir=UPLOAD_DIR):
    """Process multiple uploaded files and log them into the SQLite database."""
    if not files:
        return "No files uploaded.", get_records_df(db_path), gr.update()

    init_db(db_path)
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    uploaded_count = 0
    for file_obj in files:
        # file_obj is a Gradio file object or filepath string depending on Gradio version
        file_path = file_obj.name if hasattr(file_obj, "name") else str(file_obj)
        orig_filename = os.path.basename(file_path)
        file_size = os.path.getsize(file_path)
        file_type, _ = mimetypes.guess_type(file_path)
        if not file_type:
            file_type = "application/octet-stream"
            
        timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")

        if file_size > MAX_EMBEDDED_SIZE:
            # File > 1MB: save to uploads directory and reference relative path
            os.makedirs(upload_dir, exist_ok=True)
            dest_filename = f"{int(datetime.now(timezone.utc).timestamp())}_{orig_filename}"
            dest_path = os.path.join(upload_dir, dest_filename)
            shutil.copy(file_path, dest_path)
            
            cursor.execute("""
                INSERT INTO files (filename, file_type, file_size_bytes, upload_timestamp, is_large_file, file_content_blob, relative_path)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """, (orig_filename, file_type, file_size, timestamp, 1, None, dest_path))
        else:
            # File <= 1MB: store binary content directly in database
            with open(file_path, "rb") as f:
                blob_data = f.read()
                
            cursor.execute("""
                INSERT INTO files (filename, file_type, file_size_bytes, upload_timestamp, is_large_file, file_content_blob, relative_path)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """, (orig_filename, file_type, file_size, timestamp, 0, blob_data, None))
            
        uploaded_count += 1

    conn.commit()
    conn.close()

    df = get_records_df(db_path)
    large_files_dropdown = get_large_files_choices(db_path)
    return f"Successfully processed and logged {uploaded_count} file(s).", df, large_files_dropdown

def get_records_df(db_path=DB_NAME):
    """Retrieve database records as a list of lists for Gradio Dataframe."""
    if not os.path.exists(db_path):
        return []
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute("""
        SELECT id, filename, file_type, file_size_bytes, upload_timestamp, is_large_file, 
               CASE WHEN file_content_blob IS NOT NULL THEN 'Stored in DB (BLOB)' ELSE 'Not Stored (BLOB IS NULL)' END as blob_status,
               COALESCE(relative_path, 'N/A') as relative_path
        FROM files
        ORDER BY id DESC
    """)
    rows = cursor.fetchall()
    conn.close()
    return rows

def get_large_files_choices(db_path=DB_NAME):
    """Get list of choices for large file download dropdown."""
    if not os.path.exists(db_path):
        return gr.update(choices=[], value=None)
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute("""
        SELECT id, filename, relative_path
        FROM files
        WHERE is_large_file = 1 AND relative_path IS NOT NULL
        ORDER BY id DESC
    """)
    rows = cursor.fetchall()
    conn.close()
    
    choices = [f"ID {r[0]}: {r[1]} ({r[2]})" for r in rows]
    val = choices[0] if choices else None
    return gr.update(choices=choices, value=val)

def download_selected_large_file(selected_item, db_path=DB_NAME):
    """Return the file path for the selected large file to allow download."""
    if not selected_item:
        return None
    
    try:
        file_id = int(selected_item.split(":")[0].replace("ID", "").strip())
    except Exception:
        return None

    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute("SELECT relative_path FROM files WHERE id = ?", (file_id,))
    row = cursor.fetchone()
    conn.close()

    if row and row[0] and os.path.exists(row[0]):
        return row[0]
    return None

def download_database(db_path=DB_NAME):
    """Return database file for downloading."""
    if os.path.exists(db_path):
        return db_path
    return None

# Initialize DB on load
init_db()

# Build Gradio UI
with gr.Blocks(title="Personal SQL DB Generator") as demo:
    gr.Markdown("# 🗄️ Personal SQL DB File Generator & Logger")
    gr.Markdown(
        "Upload multiple files of any type. Files **<= 1MB** are stored directly as binary BLOBs inside the SQLite database. "
        "Files **> 1MB** are stored on disk with relative path references stored in the database. "
        "You can inspect logged entries, download the complete SQLite `.db` database file, and download big files."
    )

    with gr.Tab("Upload Files & Log to SQL"):
        file_input = gr.File(
            label="Select or Drag & Drop Multiple Files",
            file_count="multiple",
            type="filepath"
        )
        upload_button = gr.Button("Upload & Log to Database", variant="primary")
        status_output = gr.Textbox(label="Status Log", interactive=False)

    with gr.Tab("Database Records"):
        refresh_button = gr.Button("🔄 Refresh Database Table")
        db_table = gr.Dataframe(
            headers=["ID", "Filename", "Type", "Size (bytes)", "Timestamp", "Is Large (>1MB)", "BLOB Status", "Relative Path"],
            value=get_records_df(),
            interactive=False,
            wrap=True
        )

    with gr.Tab("Downloads"):
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 📦 Download SQLite Database")
                db_download_btn = gr.Button("Prepare Database File")
                db_file_output = gr.File(label="Download SQLite DB (.db)")

            with gr.Column():
                gr.Markdown("### 📂 Download Large Files (>1MB)")
                large_files_dropdown = gr.Dropdown(
                    label="Select Large File to Download",
                    choices=[c for c in get_large_files_choices()["choices"]] if isinstance(get_large_files_choices(), dict) else [],
                    interactive=True
                )
                refresh_dropdown_btn = gr.Button("Refresh Large Files List")
                download_large_btn = gr.Button("Prepare Large File Download")
                large_file_output = gr.File(label="Download File")

    # Event handlers
    upload_button.click(
        fn=process_uploaded_files,
        inputs=[file_input],
        outputs=[status_output, db_table, large_files_dropdown]
    )

    refresh_button.click(
        fn=get_records_df,
        inputs=[],
        outputs=[db_table]
    )

    db_download_btn.click(
        fn=download_database,
        inputs=[],
        outputs=[db_file_output]
    )

    refresh_dropdown_btn.click(
        fn=get_large_files_choices,
        inputs=[],
        outputs=[large_files_dropdown]
    )

    download_large_btn.click(
        fn=download_selected_large_file,
        inputs=[large_files_dropdown],
        outputs=[large_file_output]
    )

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