import marimo
import pandas as pd
import io
import matplotlib.pyplot as plt
import os
from script import getSearchResult, getClustersWithGraph, compare_clusters

app = marimo.App()

# Cell 1: Configuration Inputs
@app.cell
def cell_config(mo):
    csv_file = mo.ui.file(label="Upload Keyword CSV (1 column)")
    api_key = os.environ.get("GOOGLE_API_KEY", "")
    cse_id  = os.environ.get("GOOGLE_CSE_ID", "")
    country = mo.ui.text(label="Country Code (e.g. UK)", value="UK")
    language= mo.ui.text(label="Language Code (e.g. EN)", value="EN")
    database= mo.ui.text(label="SQLite DB Name", value="data.db")
    serp_table   = mo.ui.text(label="SERP Table",    value="keywords_serps")
    cluster_table= mo.ui.text(label="Cluster Table", value="keyword_clusters")
    timestamp_options = [
        ("Latest available (max)", "max"),
        ("March 2024 core update (2024-04-26 00:00:00.000000)", "2024-04-26 00:00:00.000000"),
        ("August 2024 core update (2024-09-10 00:00:00.000000)", "2024-09-10 00:00:00.000000"),
        ("November 2024 core update (2024-12-11 00:00:00.000000)", "2024-12-11 00:00:00.000000"),
        ("December 2024 core update (2024-12-25 00:00:00.000000)", "2024-12-25 00:00:00.000000"),
        ("March 2025 core update (2025-04-03 00:00:00.000000)", "2025-04-03 00:00:00.000000"),
        ("June 2025 core update (2025-07-24 00:00:00.000000)", "2025-07-24 00:00:00.000000"),
    ]
    ts1 = mo.ui.dropdown(timestamp_options, label="Timestamp 1")
    ts2 = mo.ui.dropdown(timestamp_options, label="Timestamp 2")
    run_btn = mo.ui.button(label="Run Comparison")
    return csv_file, api_key, cse_id, country, language, database, serp_table, cluster_table, ts1, ts2, run_btn

# Cell 2: Fetch, Cluster, Compare
@app.cell
def cell_process(csv_file, api_key, cse_id, country, language, database, serp_table, cluster_table, ts1, ts2, run_btn):
    if not (run_btn.clicked and csv_file and csv_file.value):
        return None, None, None
    # Load keywords
    df_kw = pd.read_csv(io.BytesIO(csv_file.value.read()))
    keywords = df_kw.iloc[:, 0].tolist()

    # First timestamp
    getSearchResult(keywords, language, country, api_key, cse_id, database, serp_table)
    fig1, clusters1 = getClustersWithGraph(database, serp_table, ts1.value)

    # Second timestamp
    getSearchResult(keywords, language, country, api_key, cse_id, database, serp_table)
    fig2, clusters2 = getClustersWithGraph(database, serp_table, ts2.value)

    movement = compare_clusters(clusters1, clusters2)
    return fig1, fig2, movement

# Cell 3: Display Graphs
@app.cell
def cell_display_graphs(fig1, fig2):
    if fig1 is not None and fig2 is not None:
        display(fig1)
        display(fig2)
    return None

# Cell 4: Display Movement Table
@app.cell
def cell_display_movement(movement):
    if movement is not None:
        movement = movement.copy()
        movement.sort_values(by="searchTerms", inplace=True)
        movement.reset_index(drop=True, inplace=True)
        return movement
    return None