import streamlit as st
import pandas as pd
import plotly.express as px
from sqlalchemy import create_engine, text
from sshtunnel import SSHTunnelForwarder
import os
import tempfile

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

# --- Language toggle ---
# --- Language toggle with flags ---
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(
    "",
    ["🇬🇧", "🇧🇬"],  # English and Bulgarian flags
    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 ---
SSH_KEY = st.secrets["SSH_KEY"]
SSH_USER = st.secrets["SSH_USER"]
SSH_HOST = st.secrets["DB_HOST"]
REMOTE_BIND_HOST = st.secrets["REMOTE_BIND_HOST"]
LOCAL_BIND_HOST = st.secrets["LOCAL_BIND_HOST"]
MYSQL_USER = st.secrets["MYSQL_USER"]
MYSQL_PASSWORD = st.secrets["MYSQL_PASSWORD"]
MYSQL_DB = st.secrets["DB_NAME"]
MYSQL_PORT = int(st.secrets["MYSQL_PORT"])

# --- Write SSH key to temp file securely ---
with tempfile.NamedTemporaryFile(delete=False, suffix=".pem") as temp_file:
    temp_file.write(SSH_KEY.encode())
    temp_file_path = temp_file.name
os.chmod(temp_file_path, 0o600)

# --- SSH Tunnel Setup with Error Handling ---
@st.cache_resource
def start_ssh_tunnel():
    try:
        tunnel = SSHTunnelForwarder(
            (SSH_HOST, 22),
            ssh_username=SSH_USER,
            ssh_pkey=temp_file_path,
            remote_bind_address=(REMOTE_BIND_HOST, MYSQL_PORT),
            local_bind_address=(LOCAL_BIND_HOST, MYSQL_PORT)
        )
        if not tunnel.is_active:
            tunnel.start()
        return tunnel
    except Exception as e:
        st.error(f"Failed to start SSH tunnel: {e}")
        st.stop()

try:
    server = start_ssh_tunnel()
    if not server.is_active:
        server.start()
except Exception as e:
    st.error(f"SSH tunnel error: {e}")
    st.stop()

# --- SQLAlchemy engine with Error Handling ---
try:
    MYSQL_URL = f"mysql+pymysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{LOCAL_BIND_HOST}:{MYSQL_PORT}/{MYSQL_DB}"
    engine = create_engine(MYSQL_URL)
except Exception as e:
    st.error(f"Failed to connect to the database: {e}")
    st.stop()

@st.cache_data(ttl=300)
def get_data():
    try:
        query = text("""
            SELECT id, entity, entity_score, domain, title_score, overall_score, created_at
            FROM sentiment_analysis
        """)
        df = pd.read_sql(query, engine)
        df["created_at"] = pd.to_datetime(df["created_at"], unit="s")
        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:
    filtered_df["date"] = filtered_df["created_at"].dt.floor("D")

    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()
        )

        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)
        label_avg = grouped.groupby("label")["score"].mean().reset_index(name="avg_score")
        grouped = grouped.merge(label_avg, on="label")
        grouped = grouped.sort_values(by=["avg_score", "date"], ascending=[False, True])

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

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

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

        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)
        label_avg = grouped.groupby("label")[score_type].mean().reset_index(name="avg_score")
        grouped = grouped.merge(label_avg, on="label")
        grouped = grouped.sort_values(by=["avg_score", "date"], ascending=[False, 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())
        )

    # 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",  # Set legend orientation to horizontal
            yanchor="top",
            y=-0.3,  # Position legend below the graph
            xanchor="center",
            x=0.5,
            title="",
            traceorder="normal"  # Ensure legend order matches the sorted data
        )
    )

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