import streamlit as st
import pandas as pd
import plotly.express as px
import duckdb
import os

# --- Set page layout to wide ---
st.set_page_config(layout="wide")

# --- Language toggle ---
flag_style = """
    <style>
    .flag-container {
        display: flex;
        justify-content: center;
        gap: 10px;
    }
    .flag {
        cursor: pointer;
        width: 40px;
        height: 30px;
    }
    </style>
"""
st.markdown(flag_style, unsafe_allow_html=True)

flag_selection = st.radio(
    "",
    ["🇬🇧", "🇧🇬"],
    horizontal=True,
    label_visibility="collapsed"
)
LANG = "English" if flag_selection == "🇬🇧" else "Български"

T = {
    "English": {
        "title": "📊 Sentiment Analysis Dashboard",
        "select_entities": "Select Entities",
        "select_domains": "Select Domains (News Sites)",
        "score_type": "Select Score Type",
        "group_by_domain": "Group by Domain",
        "all_scores": "Show All Score Types",
        "no_data": "No data matches the selected filters.",
        "avg_over_time": "Average {} Over Time",
        "scores_over_time": "Sentiment Scores Over Time",
        "select_entities_prompt": "Please select at least one entity to view the graph."
    },
    "Български": {
        "title": "📊 Табло за анализ на настроенията",
        "select_entities": "Изберете обекти",
        "select_domains": "Изберете източници (сайтове)",
        "score_type": "Изберете тип оценка",
        "group_by_domain": "Групирай по сайт",
        "all_scores": "Покажи всички типове оценки",
        "no_data": "Няма данни за избраните филтри.",
        "avg_over_time": "Средна стойност на {} във времето",
        "scores_over_time": "Оценки на настроенията във времето",
        "select_entities_prompt": "Моля, изберете поне един обект, за да видите графиката."
    }
}

# --- Secrets ---
if "MOTHERDUCK_TOKEN" not in st.secrets:
    st.error("MOTHERDUCK_TOKEN not found in secrets.")
    st.stop()

MOTHERDUCK_TOKEN = st.secrets["MOTHERDUCK_TOKEN"]

# --- Data Fetching (MotherDuck) ---
@st.cache_data(ttl=300)
def get_data():
    try:
        # Connect to MotherDuck using the token
        con = duckdb.connect(f'md:?token={MOTHERDUCK_TOKEN}')
        
        # Query based on your schema
        query = """
            SELECT entity, entity_score, domain, title_score, overall_score, created_at
            FROM sentiment_analysis
        """
        
        df = con.sql(query).df()
        
        # 1. Convert to Datetime (handle numeric or string inputs)
        if pd.api.types.is_numeric_dtype(df["created_at"]):
             df["created_at"] = pd.to_datetime(df["created_at"], unit="s")
        else:
             df["created_at"] = pd.to_datetime(df["created_at"])
             
        # 2. FIX: Remove Timezone Information (Make it Naive)
        if pd.api.types.is_datetime64_any_dtype(df["created_at"]):
            if df["created_at"].dt.tz is not None:
                df["created_at"] = df["created_at"].dt.tz_localize(None)

        return df
    except Exception as e:
        st.error(f"Failed to fetch data: {e}")
        return pd.DataFrame()

df = get_data()
if df.empty:
    st.warning("No data available.")
    st.stop()

# --- UI ---
st.title(T[LANG]["title"])

# --- Sidebar for filters ---
with st.sidebar:
    # Entity dropdown with counts
    entity_counts = df["entity"].value_counts()
    entity_labels = [f"{ent} ({count})" for ent, count in entity_counts.items()]
    entity_lookup = {f"{ent} ({count})": ent for ent, count in entity_counts.items()}
    selected_labels = st.multiselect(T[LANG]["select_entities"], entity_labels)
    entities = [entity_lookup[label] for label in selected_labels]

    # Domain dropdown with counts
    domain_counts = df["domain"].value_counts()
    domain_labels = [f"{dom} ({count})" for dom, count in domain_counts.items()]
    domain_lookup = {f"{dom} ({count})": dom for dom, count in domain_counts.items()}
    selected_domain_labels = st.multiselect(T[LANG]["select_domains"], domain_labels)
    domains = [domain_lookup[label] for label in selected_domain_labels]

    # Score type and grouping options
    score_type = st.selectbox(T[LANG]["score_type"], ["entity_score", "title_score", "overall_score"])
    group_by_domain = st.checkbox(T[LANG]["group_by_domain"])
    group_by_score_type = st.checkbox(T[LANG]["all_scores"])

    # Timeframe selection
    timeframes = ["All", "Last 7 Days", "Last 30 Days", "Last Year", "Custom"]
    selected_timeframe = st.selectbox("Select Timeframe", timeframes)

    # Custom date range selection
    start_date, end_date = None, None
    if selected_timeframe == "Custom":
        col1, col2 = st.columns(2)
        with col1:
            start_date = st.date_input("Start Date")
        with col2:
            end_date = st.date_input("End Date")
        if start_date > end_date:
            st.error("Start date must be before or equal to the end date.")
            st.stop()

# --- Filters ---
filtered_df = df.copy()
if entities:
    filtered_df = filtered_df[filtered_df["entity"].isin(entities)]
if domains:
    filtered_df = filtered_df[filtered_df["domain"].isin(domains)]

if selected_timeframe != "All":
    now = pd.Timestamp.now()
    if selected_timeframe == "Last 7 Days":
        start_date = now - pd.Timedelta(days=7)
        end_date = now
    elif selected_timeframe == "Last 30 Days":
        start_date = now - pd.Timedelta(days=30)
        end_date = now
    elif selected_timeframe == "Last Year":
        start_date = now - pd.Timedelta(days=365)
        end_date = now
    
    if start_date and end_date:
        filtered_df = filtered_df[
            (filtered_df["created_at"] >= pd.Timestamp(start_date)) &
            (filtered_df["created_at"] <= pd.Timestamp(end_date))
        ]

# --- Stop early if no entity selected ---
if not entities:
    st.info(T[LANG]["select_entities_prompt"])
    st.stop()

# --- Continue only if data exists ---
if not filtered_df.empty:
    
    # --- LEGEND SORTING LOGIC ---
    # We calculate the average "entity_score" for each entity/domain combo.
    # This ranking is used to sort the legend, regardless of what score is currently plotted.
    rank_group_cols = ["entity"]
    if group_by_domain:
        rank_group_cols.append("domain")
    
    # Calculate average entity_score for ranking
    rank_df = filtered_df.groupby(rank_group_cols)["entity_score"].mean().reset_index()
    # Create a dictionary for mapping: key -> score
    rank_lookup = {
        tuple(row[col] for col in rank_group_cols): row["entity_score"]
        for _, row in rank_df.iterrows()
    }

    # --- SMART DATE BINNING (Max ~20 Dots) ---
    min_date = filtered_df["created_at"].min()
    max_date = filtered_df["created_at"].max()
    days_span = (max_date - min_date).days + 1
    
    if days_span <= 20:
        freq = "D"
    else:
        step = int(days_span / 20)
        step = max(1, step)
        freq = f"{step}D"
        
    filtered_df["date"] = filtered_df["created_at"].dt.floor(freq)
    # --- END BINNING ---

    if group_by_score_type:
        score_cols = ["entity_score", "title_score", "overall_score"]
        melted = filtered_df.melt(
            id_vars=["date", "entity", "domain"],
            value_vars=score_cols,
            var_name="score_type",
            value_name="score"
        )

        group_cols = ["date", "entity"]
        if group_by_domain:
            group_cols.append("domain")
        group_cols.append("score_type")

        grouped = (
            melted.groupby(group_cols)["score"]
            .mean()
            .reset_index()
        )
        
        # FIX: Round scores to 2 decimals for the hover tooltip
        grouped["score"] = grouped["score"].round(2)

        stats_group = ["entity", "score_type"]
        if group_by_domain:
            stats_group.append("domain")

        stats = (
            melted.groupby(stats_group)["score"]
            .agg(["mean", "count"])
            .round(2)
            .reset_index()
        )

        stat_lookup = {
            tuple(row[col] for col in stats_group): f"{row['mean']:.2f}, n={int(row['count'])}"
            for _, row in stats.iterrows()
        }

        def build_label(row):
            key = tuple(row[col] for col in stats_group)
            label = f"{row['entity']}"
            if group_by_domain:
                label += f" | {row['domain']}"
            label += f" | {row['score_type'].replace('_', ' ').title()}"
            label += f" (avg: {stat_lookup.get(key, '')})"
            return label

        grouped["label"] = grouped.apply(build_label, axis=1)

        # Apply Sorting: Get rank score based on entity/domain
        def get_rank_score(row):
            key_cols = [row["entity"]]
            if group_by_domain:
                key_cols.append(row["domain"])
            return rank_lookup.get(tuple(key_cols), -999) # Default to low if missing

        grouped["rank_score"] = grouped.apply(get_rank_score, axis=1)
        
        # Sort by Rank (Entity Score), then Score Type, then Date
        grouped = grouped.sort_values(by=["rank_score", "label", "date"], ascending=[False, True, True])

        fig = px.line(
            grouped,
            x="date",
            y="score",
            color="label",
            labels={"date": "Date", "score": "Score", "label": "Legend"},
            title=T[LANG]["scores_over_time"]
        )
        fig.update_traces(mode="lines+markers")

    else:
        group_cols = ["date", "entity"]
        if group_by_domain:
            group_cols.append("domain")

        grouped = (
            filtered_df.groupby(group_cols)[score_type]
            .mean()
            .reset_index()
        )

        # FIX: Round scores to 2 decimals for the hover tooltip
        grouped[score_type] = grouped[score_type].round(2)

        stats_group = ["entity"]
        if group_by_domain:
            stats_group.append("domain")
        stats = (
            filtered_df.groupby(stats_group)[score_type]
            .agg(["mean", "count"])
            .round(2)
            .reset_index()
        )

        stat_lookup = {
            tuple(row[col] for col in stats_group): f"{row['mean']:.2f}, n={int(row['count'])}"
            for _, row in stats.iterrows()
        }

        def build_label(row):
            key = tuple(row[col] for col in stats_group)
            label = f"{row['entity']}"
            if group_by_domain:
                label += f" | {row['domain']}"
            label += f" (avg: {stat_lookup.get(key, '')})"
            return label

        grouped["label"] = grouped.apply(build_label, axis=1)
        
        # Apply Sorting: Get rank score based on entity/domain
        def get_rank_score(row):
            key_cols = [row["entity"]]
            if group_by_domain:
                key_cols.append(row["domain"])
            return rank_lookup.get(tuple(key_cols), -999)

        grouped["rank_score"] = grouped.apply(get_rank_score, axis=1)
        
        # Sort by Rank (Entity Score), then Date
        grouped = grouped.sort_values(by=["rank_score", "label", "date"], ascending=[False, True, True])

        fig = px.line(
            grouped,
            x="date",
            y=score_type,
            color="label",
            labels={"date": "Date", score_type: "Score", "label": "Legend"},
            title=T[LANG]["avg_over_time"].format(score_type.replace("_", " ").title())
        )
        fig.update_traces(mode="lines+markers")

    # Grid lines at every integer, thicker at -5 and 5
    shapes = []
    for y in range(-10, 11):
        shapes.append({
            "type": "line",
            "xref": "paper",
            "x0": 0,
            "x1": 1,
            "yref": "y",
            "y0": y,
            "y1": y,
            "line": {
                "color": "Black",
                "width": 1 if y in [-5, 5] else 0.2,
                "dash": "solid"
            }
        })

    fig.update_layout(
        xaxis=dict(type="date"),
        yaxis=dict(range=[-10, 10]),
        plot_bgcolor="white",
        shapes=shapes,
        legend=dict(
            orientation="h",
            yanchor="top",
            y=-0.3,
            xanchor="center",
            x=0.5,
            title="",
            traceorder="normal"
        )
    )

    st.plotly_chart(fig, use_container_width=True)
else:
    st.warning(T[LANG]["no_data"])