app.py

import os
import sqlite3
import io
import pandas as pd
from PIL import Image
import gradio as gr

def connect_db(db_file):
    if not db_file:
        return None
    db_path = db_file.name if hasattr(db_file, "name") else str(db_file)
    if not os.path.exists(db_path):
        return None
    return db_path

def get_all_records(db_path):
    if not db_path or not os.path.exists(db_path):
        return pd.DataFrame(), "No database loaded or file does not exist."
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='files';")
        if not cursor.fetchone():
            conn.close()
            return pd.DataFrame(), "Table 'files' not found in database."

        query = """
            SELECT id, filename, file_type, file_size_bytes, upload_timestamp, is_large_file,
                   CASE WHEN file_content_blob IS NOT NULL THEN 'BLOB Present' ELSE 'None' END as blob_status,
                   COALESCE(relative_path, 'N/A') as relative_path
            FROM files
            ORDER BY id ASC
        """
        df = pd.read_sql_query(query, conn)
        conn.close()
        return df, f"Loaded {len(df)} record(s) from database."
    except Exception as e:
        return pd.DataFrame(), f"Error reading database: {str(e)}"

def extract_file_content(row, db_dir):
    """
    Extract content (bytes or text) from a DB record tuple/dict.
    row: (id, filename, file_type, file_size_bytes, upload_timestamp, is_large_file, file_content_blob, relative_path)
    """
    filename = row['filename']
    blob = row['file_content_blob']
    rel_path = row['relative_path']

    # 1. Check BLOB content
    if blob is not None:
        return blob, filename

    # 2. Check relative_path or upload_dir fallback
    paths_to_check = []
    if rel_path:
        paths_to_check.append(rel_path)
        if db_dir:
            paths_to_check.append(os.path.join(db_dir, rel_path))
            paths_to_check.append(os.path.join(db_dir, os.path.basename(rel_path)))

    if db_dir:
        paths_to_check.append(os.path.join(db_dir, "uploads", filename))
        paths_to_check.append(os.path.join(db_dir, filename))

    for path in paths_to_check:
        if path and os.path.exists(path) and os.path.isfile(path):
            try:
                with open(path, "rb") as f:
                    return f.read(), filename
            except Exception:
                pass

    return None, filename

def is_image_file(filename, file_type=""):
    ext = os.path.splitext(filename)[1].lower()
    if ext in ['.png', '.jpg', '.jpeg', '.gif', '.webp', '.bmp', '.svg']:
        return True
    if file_type and file_type.startswith('image/'):
        return True
    return False

def is_text_file(filename, file_type=""):
    ext = os.path.splitext(filename)[1].lower()
    if ext in ['.txt', '.log', '.md', '.csv', '.json', '.yaml', '.yml']:
        return True
    if file_type and file_type.startswith('text/'):
        return True
    return False

def load_gallery_and_texts(db_file):
    db_path = connect_db(db_file)
    if not db_path:
        return [], [], "Please upload a valid SQLite database file (.db)."

    db_dir = os.path.dirname(db_path)

    try:
        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM files ORDER BY id ASC")
        rows = cursor.fetchall()
        conn.close()
    except Exception as e:
        return [], [], f"Failed to load records from DB: {str(e)}"

    gallery_images = []
    text_previews = []

    for row in rows:
        filename = row['filename']
        file_type = row['file_type'] or ""
        content_bytes, _ = extract_file_content(row, db_dir)

        if content_bytes is None:
            continue

        if is_image_file(filename, file_type):
            try:
                img = Image.open(io.BytesIO(content_bytes))
                gallery_images.append((img, filename))
            except Exception:
                pass

        if is_text_file(filename, file_type):
            try:
                text_str = content_bytes.decode('utf-8', errors='replace')
                text_previews.append(f"=== {filename} ===\n{text_str}")
            except Exception:
                pass

    joined_texts = "\n\n".join(text_previews) if text_previews else "No text (.txt) files found in DB."
    status_msg = f"Loaded {len(gallery_images)} image(s) and {len(text_previews)} text file(s)."
    return gallery_images, joined_texts, status_msg

def search_samename(db_file, search_query):
    db_path = connect_db(db_file)
    if not db_path:
        return [], "Please upload a database file first.", "No database loaded."

    if not search_query or not search_query.strip():
        return [], "Please enter a search query.", "Empty query."

    query_str = search_query.strip()
    target_basename = os.path.splitext(query_str)[0].strip().lower()

    db_dir = os.path.dirname(db_path)

    try:
        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM files")
        rows = cursor.fetchall()
        conn.close()
    except Exception as e:
        return [], f"Database query error: {str(e)}", "Error querying DB."

    matched_txt_contents = []
    matched_images = []

    for row in rows:
        filename = row['filename']
        file_type = row['file_type'] or ""
        base, ext = os.path.splitext(filename)
        base_lower = base.lower()

        # Check if basename matches search target, or full filename matches
        if base_lower == target_basename or filename.lower() == query_str.lower():
            content_bytes, _ = extract_file_content(row, db_dir)
            if content_bytes is None:
                continue

            if is_text_file(filename, file_type):
                try:
                    text_str = content_bytes.decode('utf-8', errors='replace')
                    matched_txt_contents.append(f"📄 [{filename}]\n{text_str}")
                except Exception:
                    pass

            if is_image_file(filename, file_type):
                try:
                    img = Image.open(io.BytesIO(content_bytes))
                    matched_images.append((img, filename))
                except Exception:
                    pass

    txt_output = "\n\n".join(matched_txt_contents) if matched_txt_contents else f"No .txt content matched for basename '{target_basename}'."
    status_output = f"Search for '{search_query}' finished. Found {len(matched_images)} image(s) and {len(matched_txt_contents)} text file(s)."

    return matched_images, txt_output, status_output

def refresh_all(db_file):
    df, table_status = get_all_records(connect_db(db_file))
    images, texts, gallery_status = load_gallery_and_texts(db_file)
    return images, texts, df, f"{table_status} | {gallery_status}"

# Gradio Interface
with gr.Blocks(title="Image Gallery & SQL Viewer") as demo:
    gr.Markdown("# 🖼️ Image Gallery & SQLite Database Viewer")
    gr.Markdown("Upload a SQLite database file (`.db` generated by SQLite Maker) to inspect records, view images/texts, and search samename files.")

    with gr.Row():
        db_file_input = gr.File(label="Upload SQLite Database (.db)", file_types=[".db", ".sqlite", ".sqlite3"])
        load_btn = gr.Button("Load Database", variant="primary")

    status_box = gr.Textbox(label="Status", interactive=False)

    with gr.Tabs():
        with gr.Tab("🖼️ Image Gallery & Text Viewer"):
            gallery_display = gr.Gallery(label="Loaded Images", columns=4, height="auto")
            text_viewer = gr.TextArea(label="Text Files (.txt) Content", lines=12, interactive=False)

        with gr.Tab("🔍 Samename Search"):
            gr.Markdown("Type a filename or basename (e.g., `image1` or `image1.txt`) to find the corresponding `.txt` text and same-basename image `samename.(png/jpg/etc)`.")
            with gr.Row():
                search_input = gr.Textbox(label="Search Query (e.g. sample.txt or sample)", placeholder="Enter text filename or basename...")
                search_btn = gr.Button("Search Samename", variant="primary")
            
            search_status = gr.Textbox(label="Search Status", interactive=False)
            search_text_output = gr.TextArea(label="Matched .txt File Content", lines=8, interactive=False)
            search_gallery_output = gr.Gallery(label="Matched Samename Image(s)", columns=4, height="auto")

        with gr.Tab("📊 Database Records"):
            records_table = gr.Dataframe(label="Files Table Records", interactive=False, wrap=True)

    # Event handlers
    load_btn.click(
        fn=refresh_all,
        inputs=[db_file_input],
        outputs=[gallery_display, text_viewer, records_table, status_box]
    )

    db_file_input.change(
        fn=refresh_all,
        inputs=[db_file_input],
        outputs=[gallery_display, text_viewer, records_table, status_box]
    )

    search_btn.click(
        fn=search_samename,
        inputs=[db_file_input, search_input],
        outputs=[search_gallery_output, search_text_output, search_status]
    )

    search_input.submit(
        fn=search_samename,
        inputs=[db_file_input, search_input],
        outputs=[search_gallery_output, search_text_output, search_status]
    )

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