import gradio as gr
import pandas as pd
import duckdb

PARQUET_FILES = {
    "Animagine XL 3.1": 'data.parquet',
}

preset_values = [
    "1girl, kamado nezuko, kimetsu no yaiba",
    "1girl, kanroji mitsuri, kimetsu no yaiba",
    "1girl, fern (sousou no frieren), sousou no frieren",
    "1girl, elaina (majo no tabitabi), majo no tabitabi"
]

def search_data(search_term, selected_file):
    try:
        con = duckdb.connect(database=':memory:', read_only=False)
        con.execute(f"CREATE TABLE my_table AS SELECT * FROM read_parquet('{PARQUET_FILES[selected_file]}')")
        if search_term.strip():
            results = con.execute(f"SELECT * FROM my_table WHERE lower(teks) LIKE '%{search_term.lower()}%'").fetchdf()
        else:
            search_terms_sql = " OR ".join([f"lower(teks) LIKE '%{term.lower()}%'" for term in preset_values])
            results = con.execute(f"SELECT * FROM my_table WHERE {search_terms_sql}").fetchdf()
        con.close()

        if len(results.columns) > 12:
            results = results.iloc[:, :12]
        return results
    except FileNotFoundError:
        return pd.DataFrame({'Error': ['Parquet file not found. Please check the file path.']})
    except Exception as e:
        return pd.DataFrame({'Error': [f'An error occurred: {e}']})

if __name__ == "__main__":
    with gr.Blocks() as app:
        gr.Markdown("##  Text Search for Animagine tag characters")
        with gr.Row(): # Menggunakan Row untuk layout horizontal
            with gr.Column(): # Kolom untuk input pencarian
                search_input = gr.Textbox(
                    label="Search for characters or series:",
                    placeholder="sousou no frieren",
                )
            with gr.Column(): # Kolom untuk dropdown
                file_dropdown = gr.Dropdown(
                    choices=list(PARQUET_FILES.keys()),
                    label="Select Character Tag",
                    value=list(PARQUET_FILES.keys())[0], # Default value
                )
        search_output = gr.Dataframe( # Output di luar row, di bawahnya
            label="Search Results",
            value=pd.DataFrame({'Characters tag': preset_values}),
            headers="auto",
        )

        inputs = [search_input, file_dropdown]
        search_input.change(
            fn=search_data,
            inputs=inputs,
            outputs=search_output,
        )
        file_dropdown.change(
            fn=search_data,
            inputs=inputs,
            outputs=search_output,
        )

    app.launch()