"""
Worker script to run the SCIP solver in a separate process.
This allows capturing the C++ stdout (EnableOutput) and streaming it back to Streamlit.
"""

import sys
import pandas as pd
import math
import datetime
from ortools.linear_solver import pywraplp
import json
from vincenty import vincenty

import os
import sqlite3

# Ensure src module is visible if run standalone
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from src.fleet_assignment.config import DB_PATH, RESULTS_FILE, COST_FILE

def run_solver(time_limit_sec=300):
    print(f"--- Starting Global Solver Worker (Limit: {time_limit_sec}s) ---")
    JSON_AEROPORTOS = "airports.json"
    
    conn = sqlite3.connect(DB_PATH)
    df_frota_db = pd.read_sql_query("SELECT * FROM frota", conn)
    demands_db = pd.read_sql_query("SELECT data_apenas as 'Data Apenas', local_dec as 'Localidade Decolagem', local_pou as 'Localidade Pouso', pax as 'PAX' FROM demandas", conn)
    conn.close()
    
    df_aeroportos = pd.read_json(JSON_AEROPORTOS, orient='index')
    
    # Format demands as required by downstream logic
    pax_demanda_completa = demands_db.groupby(['Data Apenas', 'Localidade Decolagem', 'Localidade Pouso'])['PAX'].sum().reset_index()
    demands = pax_demanda_completa[pax_demanda_completa['PAX'] > 0].copy()

    with open('emissores.json', 'r') as f:
        emissores_data = json.load(f)[0]
    icaos_emissores = list(emissores_data.values())
    
    icaos_dec = demands['Localidade Decolagem'].unique().tolist()
    icaos_pou = demands['Localidade Pouso'].unique().tolist()
    all_icaos = [icao for icao in list(set(icaos_dec + icaos_pou + icaos_emissores)) if str(icao).strip() not in ['0', '']]
    
    coords = df_aeroportos.loc[df_aeroportos.index.intersection(all_icaos), ['lat', 'lon']]
    valid_icaos = coords.index.tolist()
    
    dist_matrix = pd.DataFrame(index=valid_icaos, columns=valid_icaos)
    for i in valid_icaos:
        for j in valid_icaos:
            if i == j: dist_matrix.loc[i, j] = 0.0
            else:
                lat1, lon1 = coords.loc[i, 'lat'], coords.loc[i, 'lon']
                lat2, lon2 = coords.loc[j, 'lat'], coords.loc[j, 'lon']
                dist_matrix.loc[i, j] = round(vincenty((lat1, lon1), (lat2, lon2)), 2)

    AERONAVES_INTERESSE = ['C-97', 'C-95M', 'C-105', 'KC-390', 'KC-30', 'C-99A', 'C-98', 'C-98A']
    CAPACIDADES_AERONAVES_INTERESSE = [30, 12, 73, 80, 238, 50, 10, 14]
    ALCANCE_AERONAVES_INTERESSE = [1600, 1900, 5000, 6000, 14500, 2200, 2400, 2400]
    DIARIAS_AERONAVES_INTERESSE = [x * 60 for x in [4, 3, 3, 3, 10, 3, 3, 3]]

    frotas = []
    for idx, row in df_frota_db.iterrows():
        modelo, emissor = row['Modelo'], row['Emissor']
        consumo = row['Consumo (km/l)']
        if consumo > 0 and modelo in AERONAVES_INTERESSE:
            a_idx = AERONAVES_INTERESSE.index(modelo)
            frotas.append({
                'modelo': str(modelo),
                'emissor': str(emissor),
                'base': emissores_data.get(str(emissor), None),
                'frota_max': int(row['Tamanho da Frota']),
                'consumo_kml': float(consumo),
                'capacidade': int(CAPACIDADES_AERONAVES_INTERESSE[a_idx]),
                'alcance': float(ALCANCE_AERONAVES_INTERESSE[a_idx]),
                'velocidade_media': float(row['Velocidade Média (km/h)']),
                'diaria_litros': float(DIARIAS_AERONAVES_INTERESSE[a_idx])
            })
    frotas = [f for f in frotas if f['base'] is not None and f['base'] in valid_icaos]

    df_dias = pd.to_datetime(demands['Data Apenas']).dt.date
    dias_unicos_list = sorted(df_dias.dropna().unique())
    if not dias_unicos_list:
        print("NO DEMAND DATA.")
        return
        
    min_date = min(dias_unicos_list)
    max_date = max(dias_unicos_list)
    todas_as_datas = [min_date + datetime.timedelta(days=i) for i in range((max_date - min_date).days + 1)]
    
    print("Pre-computing distances and overnight stays...")
    unique_routes = demands[['Localidade Decolagem', 'Localidade Pouso']].drop_duplicates()
    route_info = {}
    for idx, row in unique_routes.iterrows():
        l_dec = str(row['Localidade Decolagem'])
        l_pou = str(row['Localidade Pouso'])
        if l_dec not in valid_icaos or l_pou not in valid_icaos:
            continue
        for f in frotas:
            m, e, base = f['modelo'], f['emissor'], f['base']
            dist_total = dist_matrix.loc[base, l_dec] + dist_matrix.loc[l_dec, l_pou] + dist_matrix.loc[l_pou, base]
            tempo_missao = dist_total / f['velocidade_media'] if f['velocidade_media'] > 0 else float('inf')
            
            if tempo_missao <= 96.0:
                num_pernoites = max(0, math.ceil(tempo_missao / 12.0) - 1)
                fator1 = 1.25 if dist_matrix.loc[base, l_dec] > f['alcance'] else 1.0
                fator2 = 1.25 if dist_matrix.loc[l_dec, l_pou] > f['alcance'] else 1.0
                fator3 = 1.25 if dist_matrix.loc[l_pou, base] > f['alcance'] else 1.0
                comb_missao = ((dist_matrix.loc[base, l_dec]*fator1 + dist_matrix.loc[l_dec, l_pou]*fator2 + dist_matrix.loc[l_pou, base]*fator3) / f['consumo_kml']) + (num_pernoites * f['diaria_litros'])
                route_info[(m, e, l_dec, l_pou)] = {'num_pernoites': num_pernoites, 'combustivel_missao': comb_missao}

    solver = pywraplp.Solver.CreateSolver('SCIP')
    solver.SetTimeLimit(time_limit_sec * 1000)
    solver.EnableOutput() # O MAIS IMPORTANTE: Isso escreve no stdout para o subprocess ler!
    
    print("Building Demand Matrix...")
    x, s = {}, {}
    for t in dias_unicos_list:
        demandas_dia = demands[df_dias == t]
        for idx, d_row in demandas_dia.iterrows():
            l_dec, l_pou = str(d_row['Localidade Decolagem']), str(d_row['Localidade Pouso'])
            if l_dec not in valid_icaos or l_pou not in valid_icaos: continue
            
            s[(l_dec, l_pou, t)] = solver.NumVar(0, solver.infinity(), f"s_{l_dec}_{l_pou}_{t}")
            restricao_pax = solver.Constraint(float(d_row['PAX']), solver.infinity(), "")
            restricao_pax.SetCoefficient(s[(l_dec, l_pou, t)], 1)
            
            for f in frotas:
                m, e = f['modelo'], f['emissor']
                if (m, e, l_dec, l_pou) in route_info:
                    x[(m, e, l_dec, l_pou, t)] = solver.IntVar(0, solver.infinity(), f"x_{m}_{e}_{l_dec}_{l_pou}_{t}")
                    restricao_pax.SetCoefficient(x[(m, e, l_dec, l_pou, t)], f['capacidade'])

    print("Building Fleet Temporal Matrix...")
    for f in frotas:
        m, e, frota_max = f['modelo'], f['emissor'], f['frota_max']
        for t_atual in todas_as_datas:
            restricao_frota = solver.Constraint(0, frota_max, "")
            for t_start in dias_unicos_list:
                diff_days = (t_atual - t_start).days
                if 0 <= diff_days <= 7:
                    for idx, d_row in demands[df_dias == t_start].iterrows():
                        l_dec, l_pou = str(d_row['Localidade Decolagem']), str(d_row['Localidade Pouso'])
                        if (m, e, l_dec, l_pou, t_start) in x:
                            if diff_days <= route_info[(m, e, l_dec, l_pou)]['num_pernoites']:
                                restricao_frota.SetCoefficient(x[(m, e, l_dec, l_pou, t_start)], 1)

    print("\n--- STARTING SCIP OPTIMIZATION ---\n")
    sys.stdout.flush()
    
    objetivo = solver.Objective()
    objetivo.SetMinimization()
    for var in s.values(): objetivo.SetCoefficient(var, 1e8)
    for (m, e, l_dec, l_pou, t), var in x.items():
        objetivo.SetCoefficient(var, route_info[(m, e, l_dec, l_pou)]['combustivel_missao'])
        
    status = solver.Solve()
    
    print(f"\n--- SCIP FINISHED ---")
    if status in [pywraplp.Solver.OPTIMAL, pywraplp.Solver.FEASIBLE]:
        obj_val = solver.Objective().Value()
        try:
            best_bound = solver.Objective().BestBound()
        except:
            best_bound = obj_val
        gap = abs((obj_val - best_bound) / obj_val) * 100 if obj_val > 0 else 0
        
        print(f"Final Status: {'OPTIMAL' if status == pywraplp.Solver.OPTIMAL else 'FEASIBLE (Time Limit)'}")
        print(f"Optimized Cost: {obj_val:,.2f} L")
        with open(COST_FILE, "w") as f:
            f.write(str(obj_val))
        print(f"Dual Bound: {best_bound:,.2f} L")
        print(f"Gap: {gap:.4f}%")
        
        results = []
        for (m, e, l_dec, l_pou, t), var in x.items():
            if var.solution_value() > 0:
                results.append({
                    'Data': t, 'Modelo': m, 'Emissor': e,
                    'Local Decolagem': l_dec, 'Local Pouso': l_pou,
                    'Qtd Aeronaves': int(var.solution_value()),
                    'Pernoites': route_info[(m, e, l_dec, l_pou)]['num_pernoites']
                })
        df_results = pd.DataFrame(results)
        df_results.to_csv(RESULTS_FILE, sep=';', index=False)
        print(f"Results successfully saved to {RESULTS_FILE}")
    else:
        print("\n[ERROR] Solver failed to find a solution.")

if __name__ == '__main__':
    time_limit = int(sys.argv[1]) if len(sys.argv) > 1 else 300
    run_solver(time_limit)
