"""
Streamlit Dashboard for Brazilian Air Force Fleet Assignment.
Provides interactive mapping, scenario generation, and SCIP solver integration.
"""

import os
import json
import sqlite3
import subprocess
import datetime
import pandas as pd
import pydeck as pdk
import streamlit as st
from vincenty import vincenty

# ==========================================
# PAGE CONFIGURATION
# ==========================================
st.set_page_config(
    page_title="Fleet Assignment",
    page_icon="✈️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ==========================================
# CONSTANTS & CONFIGURATION
# ==========================================
import sys
# Ensure we can import from src
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.fleet_assignment.config import (
    CSV_FILEPATH, JSON_AEROPORTOS, JSON_EMISSORES, DB_PATH,
    RESULTS_FILE, COST_FILE, AIRCRAFT_CONFIG
)

from src.fleet_assignment.ingest import load_data

@st.cache_data
def get_cached_data():
    return load_data()

demands_raw, df_fleet, valid_icaos, airports_geo = get_cached_data()

# ==========================================
# DATABASE MANAGEMENT
# ==========================================
def execute_sql(query, params=(), fetch=False):
    """Helper to execute SQL commands cleanly."""
    with sqlite3.connect(DB_PATH) as conn:
        if fetch:
            return pd.read_sql_query(query, conn, params=params)
        conn.execute(query, params)
        conn.commit()

def init_db(demands_df, frota_df, force_reset=False):
    """Initializes SQLite Database as the single source of truth."""
    if force_reset and os.path.exists(DB_PATH):
        os.remove(DB_PATH)
    
    with sqlite3.connect(DB_PATH) as conn:
        conn.execute('''
            CREATE TABLE IF NOT EXISTS demandas (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                data_apenas TEXT,
                local_dec TEXT,
                local_pou TEXT,
                pax INTEGER
            )
        ''')
        
        count = conn.execute("SELECT COUNT(*) FROM demandas").fetchone()[0]
        if count == 0:
            df_insert = demands_df[['Data Apenas', 'Localidade Decolagem', 'Localidade Pouso', 'PAX']].copy()
            df_insert.columns = ['data_apenas', 'local_dec', 'local_pou', 'pax']
            df_insert.to_sql('demandas', conn, if_exists='append', index=False)
            frota_df.to_sql('frota', conn, if_exists='replace', index=False)

def load_demands():
    return execute_sql("SELECT id as ID, data_apenas as 'Date', local_dec as 'Dep', local_pou as 'Arr', pax as 'PAX' FROM demandas", fetch=True)

init_db(demands_raw, df_fleet)

@st.cache_data
def load_results():
    try:
        return pd.read_csv(RESULTS_FILE, sep=';')
    except FileNotFoundError:
        return pd.DataFrame()

df_results = load_results()

# ==========================================
# SOLVER INTEGRATION
# ==========================================
def run_solver_subprocess(time_limit_sec, output_placeholder):
    """Spawns the optimization worker and streams logs to the UI."""
    process = subprocess.Popen(
        ["uv", "run", "python", "-u", "solver_worker.py", str(time_limit_sec)],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1
    )
    
    full_log = ""
    for line in iter(process.stdout.readline, ''):
        full_log += line
        output_placeholder.code(full_log, language='text')
        
    process.wait()
    return process.returncode == 0

# ==========================================
# UI RENDER
# ==========================================
st.title("✈️ Fleet Assignment")
st.markdown("Brazilian Air Force Fleet Routing and Scheduling (Global Time-Space Network Flow)")

tab1, tab2, tab3, tab4, tab5 = st.tabs(["Demand Data", "Fleet", "Solver", "Map", "Results"])

# --- TAB 1: DEMAND DATA ---
with tab1:
    col_title, col_reset = st.columns([0.8, 0.2])
    with col_title:
        st.header("Passenger Demands")
    with col_reset:
        if st.button("🔄 Reset DB to Raw CSV", type="secondary", width="stretch"):
            init_db(demands_raw, df_fleet, force_reset=True)
            st.success("Database Reset!")
            st.rerun()
            
    df_display_demand = load_demands()
    
    st.info("💡 You can edit, add, or delete rows directly in the table below. Click 'Save Changes' when done.")
    edited_demand = st.data_editor(
        df_display_demand,
        num_rows="dynamic",
        width="stretch",
        hide_index=True,
        key="demand_editor"
    )
    
    if st.button("💾 Save Demand Changes", type="primary", width="stretch"):
        with sqlite3.connect(DB_PATH) as conn:
            conn.execute("DROP TABLE IF EXISTS demandas")
            conn.execute('''
                CREATE TABLE demandas (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    data_apenas TEXT,
                    local_dec TEXT,
                    local_pou TEXT,
                    pax INTEGER
                )
            ''')
            df_to_save = edited_demand.drop(columns=['ID'], errors='ignore')
            df_to_save.columns = ['data_apenas', 'local_dec', 'local_pou', 'pax']
            df_to_save.to_sql('demandas', conn, if_exists='append', index=False)
        st.success("Database synchronized successfully!")
        st.rerun()

# --- TAB 2: FLEET ---
with tab2:
    st.header("Fleet Specifications")
    
    # Check if we have modified fleet in DB, otherwise use default
    with sqlite3.connect(DB_PATH) as conn:
        try:
            current_fleet = pd.read_sql_query("SELECT * FROM frota", conn)
        except Exception:
            current_fleet = df_fleet.copy()
            
    df_display_fleet = current_fleet[['Modelo', 'Emissor', 'Tamanho da Frota', 'Consumo (km/l)', 'Velocidade Média (km/h)', 'Capacidade', 'Alcance']].copy()
    df_display_fleet.columns = ['Modelo', 'Squadron', 'Fleet Size', 'Fuel Usage (km/l)', 'Mean Speed (km/h)', 'PAX Capacity', 'Max Range (km)']
    
    st.info("💡 Edit specs directly below to simulate larger fleets, different fuel burn, or cargo capacity.")
    edited_fleet = st.data_editor(
        df_display_fleet,
        width="stretch",
        hide_index=True,
        key="fleet_editor"
    )
    
    if st.button("💾 Save Fleet Changes", type="primary", width="stretch"):
        rename_back = {
            'Modelo': 'Modelo', 'Squadron': 'Emissor', 'Fleet Size': 'Tamanho da Frota',
            'Fuel Usage (km/l)': 'Consumo (km/l)', 'Mean Speed (km/h)': 'Velocidade Média (km/h)',
            'PAX Capacity': 'Capacidade', 'Max Range (km)': 'Alcance'
        }
        df_save_fleet = edited_fleet.rename(columns=rename_back)
        
        with sqlite3.connect(DB_PATH) as conn:
            old_frota = pd.read_sql_query("SELECT * FROM frota", conn)
            for idx, row in df_save_fleet.iterrows():
                mask = (old_frota['Modelo'] == row['Modelo']) & (old_frota['Emissor'] == row['Emissor'])
                old_frota.loc[mask, 'Tamanho da Frota'] = row['Tamanho da Frota']
                old_frota.loc[mask, 'Consumo (km/l)'] = row['Consumo (km/l)']
                old_frota.loc[mask, 'Velocidade Média (km/h)'] = row['Velocidade Média (km/h)']
                old_frota.loc[mask, 'Capacidade'] = row['Capacidade']
                old_frota.loc[mask, 'Alcance'] = row['Alcance']
            old_frota.to_sql('frota', conn, if_exists='replace', index=False)
            
        st.success("Fleet changes saved successfully!")
        st.rerun()

# --- TAB 3: SOLVER ---
with tab3:
    st.header("Mathematical Formulation")
    st.latex(r"\min \sum_{t} \sum_{r} (c_{r} \cdot x_{r,t}) + \sum_{t} \sum_{d} (M \cdot s_{d,t})")
    st.latex(r"\text{s.t.} \quad \sum_{m,e} (\text{cap}_{m} \cdot x_{m,e,r,t}) + s_{r,t} \ge \text{PAX}_{r,t} \quad \forall r, t")
    st.latex(r"\sum_{t_{start} \le T \le t_{start} + \text{pernoites}} x_{m,e,r,t_{start}} \le \text{MaxFleet}_{m,e} \quad \forall m, e, T")
    
    st.markdown(r"""
    **Dictionary of Variables:**
    - $x_{m,e,r,t}$: Decision Variable (Integer). Number of aircraft of model $m$, from squadron $e$, allocated to route $r$ on day $t$.
    - $c_r$: Estimated total fuel cost for the route $r$ mission, including overnight stay penalties if applicable.
    - $s_{r,t}$: Slack Variable. Represents the passenger demand that could **not** be met on that day due to fleet limitations.
    - $M$: High Penalty (Big M). Imposed on the system for each unit of slack activated, forcing the solver to meet demand whenever possible.
    - $\text{cap}_m$: Maximum passenger capacity (seats) of aircraft $m$.
    - $\text{PAX}_{r,t}$: Actual number of Air Force passengers needing to fly on route $r$ on day $t$.
    - $\text{MaxFleet}_{m,e}$: Total physical aircraft available in squadron $e$ for model $m$.
    - $T$: Temporal inspection window. Ensures aircraft blockade during overnight stays, preventing fleet duplication.
    """)
    
    if st.button("🚀 Run Global Optimization", type="primary"):
        st.info("Starting Worker... Follow the real-time SCIP log below:")
        log_placeholder = st.empty()
        
        if run_solver_subprocess(300, log_placeholder):
            st.success("Solver Finished! Results saved to CSV.")
            st.session_state['new_results'] = True
        else:
            st.error("Solver Failed or was aborted.")

# --- TAB 4: MAP ---
with tab4:
    st.header("Interactive Flight Map")
    
    if st.session_state.get('new_results'):
        df_results = load_results()
        
    view_state = pdk.ViewState(latitude=-15.78, longitude=-47.92, zoom=3.5, pitch=45)
    layers = []
        
    if df_results.empty:
        st.info("No optimization results yet. Run the Global Optimization to see routes.")
    else:
        df_results['Data'] = pd.to_datetime(df_results['Data'])
        dias_disponiveis = sorted(df_results['Data'].dt.date.unique())
        
        # Safe initialization of map date
        if 'current_map_date' not in st.session_state or st.session_state.current_map_date not in dias_disponiveis:
             st.session_state.current_map_date = dias_disponiveis[0]
             
        current_idx = dias_disponiveis.index(st.session_state.current_map_date)
        
        # Navigation Controls
        col_prev, col_drop, col_next = st.columns([1, 8, 1])
        with col_prev:
            st.markdown("<br/>", unsafe_allow_html=True)
            if st.button("⬅️", disabled=current_idx==0, width="stretch"):
                st.session_state.current_map_date = dias_disponiveis[current_idx - 1]
                st.rerun()
        with col_drop:
            selected_date = st.date_input("Select via Calendar", value=st.session_state.current_map_date)
            if selected_date != st.session_state.current_map_date:
                st.session_state.current_map_date = selected_date
                st.rerun()
        with col_next:
            st.markdown("<br/>", unsafe_allow_html=True)
            if st.button("➡️", disabled=current_idx>=len(dias_disponiveis)-1, width="stretch"):
                st.session_state.current_map_date = dias_disponiveis[current_idx + 1]
                st.rerun()
                
        slider_date = st.select_slider("Or scrub through scheduled flight days", options=dias_disponiveis, value=st.session_state.current_map_date)
        if slider_date != st.session_state.current_map_date:
            st.session_state.current_map_date = slider_date
            st.rerun()
                
        df_day = df_results[df_results['Data'].dt.date == st.session_state.current_map_date]
        df_demands_map = load_demands()
        
        # Map Data Parsing
        map_data = []
        for _, row in df_day.iterrows():
            orig, dest = row['Local Decolagem'], row['Local Pouso']
            
            if orig in airports_geo and dest in airports_geo:
                lon1, lat1 = float(airports_geo[orig]['lon']), float(airports_geo[orig]['lat'])
                lon2, lat2 = float(airports_geo[dest]['lon']), float(airports_geo[dest]['lat'])
                
                pernoites = row.get('Pernoites', 0)
                color = [255, 65, 54, 200] if pernoites > 0 else [0, 116, 217, 200]
                
                # Calculate PAX sum for this route on this day
                pax_mask = (df_demands_map['Dep'] == orig) & (df_demands_map['Arr'] == dest) & (df_demands_map['Date'] == str(st.session_state.current_map_date))
                pax_count = df_demands_map[pax_mask]['PAX'].sum()
                
                map_data.append({
                    "start": [lon1, lat1], "end": [lon2, lat2], "color": color,
                    "modelo": row['Modelo'], "orig": orig, "dest": dest,
                    "pax": int(pax_count), "pernoites": pernoites
                })
        
        df_map = pd.DataFrame(map_data)
        if df_map.empty:
            st.info("No flights scheduled for this day.")
        else:
            layers.append(pdk.Layer(
                "ArcLayer", data=df_map,
                get_source_position="start", get_target_position="end",
                get_source_color="color", get_target_color="color",
                get_width=3, pickable=True
            ))
            
    st.pydeck_chart(pdk.Deck(
        layers=layers,
        initial_view_state=view_state,
        tooltip={
            "html": "<b>✈️ Aircraft:</b> {modelo}<br/><b>📍 Route:</b> {orig} ➔ {dest}<br/><b>👥 PAX:</b> {pax}<br/><b>🌙 Out of Base:</b> {pernoites} days",
            "style": {"backgroundColor": "steelblue", "color": "white"}
        } if layers else None,
        map_style='road'
    ))
    
    # Selected Day Results Table
    if not df_results.empty:
        df_day_display = df_day.copy()
        
        # Bind PAX back into the table
        df_day_display['PAX'] = df_day_display.apply(lambda r: df_demands_map[
            (df_demands_map['Dep'] == r['Local Decolagem']) & 
            (df_demands_map['Arr'] == r['Local Pouso']) & 
            (df_demands_map['Date'] == str(st.session_state.current_map_date))
        ]['PAX'].sum(), axis=1)
        
        df_day_display.rename(columns={
            'Data': 'Date', 'Modelo': 'Model', 'Emissor': 'Squadron',
            'Local Decolagem': 'Departure', 'Local Pouso': 'Arrival',
            'Qtd Aeronaves': 'Allocated Fleet', 'Pernoites': 'Days Out of Base'
        }, inplace=True)
        
        st.subheader(f"Flight Schedule for {st.session_state.current_map_date}")
        st.dataframe(df_day_display, width="stretch")

# --- TAB 5: RESULTS ---
with tab5:
    st.header("Optimization Results")
    
    try:
        raw_fuel_df = execute_sql("SELECT SUM(\"Combustível\") as total FROM frota", fetch=True)
        raw_fuel = raw_fuel_df['total'].iloc[0] if not raw_fuel_df.empty else 0
        if pd.isna(raw_fuel): raw_fuel = 0
        
        opt_fuel = None
        if os.path.exists(COST_FILE):
            with open(COST_FILE, "r") as f:
                opt_fuel = float(f.read().strip())
                
        colA, colB, colC = st.columns(3)
        with colA:
            st.metric(label="Baseline Fuel (Non-Optimized)", value=f"{raw_fuel:,.2f} L")
        with colB:
            if opt_fuel is not None:
                st.metric(label="Optimized Fuel Consumption", value=f"{opt_fuel:,.2f} L", delta=f"{opt_fuel - raw_fuel:,.2f} L", delta_color="inverse")
        with colC:
            if opt_fuel is not None and raw_fuel > 0:
                savings = (1 - (opt_fuel / raw_fuel)) * 100
                st.metric(label="Total Fuel Savings", value=f"{savings:.2f}%")
    except Exception:
        pass
        
    if not df_results.empty:
        st.dataframe(df_results, width="stretch")
        
        if 'Pernoites' in df_results.columns:
            st.subheader("Days Out of Base (Pernoites) Breakdown")
            df_pernoites = df_results[df_results['Pernoites'] > 0]
            st.dataframe(df_pernoites[['Modelo', 'Local Decolagem', 'Local Pouso', 'Pernoites']].sort_values(by='Pernoites', ascending=False))
    else:
        st.warning("No results available.")
