import duckdb
import pandas as pd

from config import DATASET_ID, HF_PARQUET_URL_1, HF_PARQUET_URL_2


def create_fresh_duckdb_with_views(parquet_url_1: str = HF_PARQUET_URL_1, parquet_url_2: str = HF_PARQUET_URL_2):
    """Return a fresh DuckDB connection with parquet-backed views configured."""
    local_con = duckdb.connect(database=":memory:", read_only=False)
    try:
        try:
            local_con.execute("INSTALL httpfs;")
            local_con.execute("LOAD httpfs;")
        except Exception:
            pass

        try:
            local_con.execute("SET enable_http_metadata_cache = false;")
            local_con.execute("SET enable_object_cache = false;")
        except Exception:
            pass

        local_con.execute(
            f"""
            CREATE OR REPLACE VIEW all_downloads AS
            SELECT * FROM read_parquet('{parquet_url_1}')
        """
        )
        local_con.execute(
            f"""
            CREATE OR REPLACE VIEW one_year_rolling AS
            SELECT * FROM read_parquet('{parquet_url_2}')
        """
        )
    except Exception:
        pass
    return local_con


def get_last_updated():
    """Return the latest timestamp available in the all_downloads view."""
    try:
        conn = create_fresh_duckdb_with_views()
        try:
            result = conn.execute("SELECT MAX(time) as max_time FROM all_downloads").fetchdf()
        finally:
            conn.close()

        max_time = result["max_time"].iloc[0]
        if pd.isnull(max_time):
            return "N/A"
        dt = pd.to_datetime(max_time)
        return dt.strftime("%b %d, %Y")
    except Exception:
        return "N/A"


def build_leaderboard_query(
    group_col,
    top_n,
    start_str=None,
    end_str=None,
    date_str=None,
    view="all_downloads",
):
    """Build the SQL query string for leaderboard data."""
    is_alltime = date_str is not None

    if group_col == "org_country_single":
        group_expr = """CASE 
            WHEN org_country_single IN ('HF', 'United States of America') THEN 'United States of America'
            WHEN org_country_single IN ('International', 'Online', 'Online?') THEN 'International/Online'
            ELSE org_country_single
        END"""
    else:
        group_expr = group_col

    if is_alltime:
        base_where = f"WHERE time <= '{date_str}'"
        downloads_calc = f"COALESCE(MAX(CASE WHEN time <= '{date_str}' THEN downloadsAllTime END), 0) AS total_downloads"
    else:
        base_where = ""
        downloads_calc = f"""COALESCE(MAX(CASE WHEN time <= '{end_str}' THEN downloadsAllTime END), 0)
            - COALESCE(MAX(CASE WHEN time < '{start_str}' THEN downloadsAllTime END), 0)
            AS total_downloads"""

    if group_col == "derived_author":
        return f"""
        WITH base_data AS (
            SELECT
                {group_expr} AS group_key,
                CASE 
                    WHEN org_country_single IN ('HF', 'United States of America') THEN 'United States of America'
                    WHEN org_country_single IN ('International', 'Online', 'Online?') THEN 'International/Online'
                    ELSE org_country_single
                END AS org_country_single,
                author,
                derived_author,
                merged_country_groups_single,
                merged_modality,
                model,
                time,
                downloadsAllTime
            FROM {view}
            {base_where}
        ),

        author_country_lookup AS (
            SELECT DISTINCT
                derived_author,
                FIRST_VALUE(org_country_single) OVER (PARTITION BY derived_author ORDER BY downloadsAllTime DESC) AS derived_author_country
            FROM base_data
            WHERE derived_author IS NOT NULL
        ),

        author_merged_country_lookup AS (
            SELECT DISTINCT
                derived_author,
                FIRST_VALUE(merged_country_groups_single) OVER (PARTITION BY derived_author ORDER BY downloadsAllTime DESC) AS derived_author_merged_country
            FROM base_data
            WHERE derived_author IS NOT NULL
        ),

        model_metrics AS (
            SELECT
                model,
                group_key,
                ANY_VALUE(org_country_single) AS org_country_single,
                ANY_VALUE(author) AS author,
                ANY_VALUE(derived_author) AS derived_author,
                ANY_VALUE(merged_country_groups_single) AS merged_country_groups_single,
                ANY_VALUE(merged_modality) AS merged_modality,
                {downloads_calc}
            FROM base_data
            GROUP BY model, group_key
        ),

        total_downloads_cte AS (
            SELECT SUM(total_downloads) AS total_downloads_all FROM model_metrics
        )

        SELECT
            mm.model,
            mm.group_key,
            acl.derived_author_country AS org_country_single,
            amc.derived_author_merged_country AS merged_country_groups_single,
            mm.author,
            mm.derived_author,
            mm.merged_modality,
            mm.total_downloads,
            CASE WHEN td.total_downloads_all = 0 THEN 0 ELSE ROUND(mm.total_downloads * 100.0 / td.total_downloads_all, 2) END AS percent_of_total
        FROM model_metrics mm
        LEFT JOIN author_country_lookup acl ON mm.group_key = acl.derived_author
        LEFT JOIN author_merged_country_lookup amc ON mm.group_key = amc.derived_author
        CROSS JOIN total_downloads_cte td
        WHERE mm.total_downloads > 0
        ORDER BY mm.total_downloads DESC
        LIMIT {top_n * 10};
        """

    return f"""
    WITH base_data AS (
        SELECT
            {group_expr} AS group_key,
            CASE 
                WHEN org_country_single IN ('HF', 'United States of America') THEN 'United States of America'
                WHEN org_country_single IN ('International', 'Online', 'Online?') THEN 'International/Online'
                ELSE org_country_single
            END AS org_country_single,
            author,
            derived_author,
            merged_country_groups_single,
            merged_modality,
            model,
            time,
            downloadsAllTime
        FROM {view}
        {base_where}
    ),

    model_metrics AS (
        SELECT
            model,
            group_key,
            ANY_VALUE(org_country_single) AS org_country_single,
            ANY_VALUE(author) AS author,
            ANY_VALUE(derived_author) AS derived_author,
            ANY_VALUE(merged_country_groups_single) AS merged_country_groups_single,
            ANY_VALUE(merged_modality) AS merged_modality,
            {downloads_calc}
        FROM base_data
        GROUP BY model, group_key
    ),

    total_downloads_cte AS (
        SELECT SUM(total_downloads) AS total_downloads_all FROM model_metrics
    )

    SELECT
        mm.model,
        mm.group_key,
        mm.org_country_single,
        mm.author,
        mm.derived_author,
        mm.merged_country_groups_single,
        mm.merged_modality,
        mm.total_downloads,
        CASE WHEN td.total_downloads_all = 0 THEN 0 ELSE ROUND(mm.total_downloads * 100.0 / td.total_downloads_all, 2) END AS percent_of_total
    FROM model_metrics mm
    CROSS JOIN total_downloads_cte td
    WHERE mm.total_downloads > 0
    ORDER BY mm.total_downloads DESC
    LIMIT {top_n * 10};
    """


def get_top_n_from_duckdb(con, group_col, top_n=10, time_filter=None, view="all_downloads"):
    """Query DuckDB directly to get model-level rows with per-model total_downloads."""
    if time_filter and len(time_filter) == 2:
        start = pd.to_datetime(time_filter[0], unit="s")
        end = pd.to_datetime(time_filter[1], unit="s")
    else:
        start = pd.to_datetime("1970-01-01")
        end = pd.Timestamp.now()

    start_str = str(start)
    end_str = str(end)
    query = build_leaderboard_query(group_col, top_n, start_str, end_str, view=view)

    conn_local = create_fresh_duckdb_with_views()
    try:
        return conn_local.execute(query).fetchdf()
    except Exception as exc:
        print(f"Error querying DuckDB: {exc}")
        return pd.DataFrame()
    finally:
        conn_local.close()