import json
import pandas as pd
from vincenty import vincenty
import sys
import os

# 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 (
    CSV_FILEPATH, JSON_AEROPORTOS, JSON_EMISSORES, AIRCRAFT_CONFIG
)

def load_data():
    """Reads raw CSV and JSON to build the baseline demand and fleet dataframes."""
    df = pd.read_csv(CSV_FILEPATH, low_memory=False)
    df_airports = pd.read_json(JSON_AEROPORTOS, orient='index')
    
    # Basic Cleaning
    df.replace('NIL', 0, inplace=True)
    df = df[(df['Localidade Decolagem'] != 0) & (df['Localidade Pouso'] != 0)]
    df = df[df['Modelo'].isin(AIRCRAFT_CONFIG.keys())]
    
    df['PAX'] = pd.to_numeric(df['PAX']).fillna(0)
    df['Dep TimeStamp'] = pd.to_datetime(df['Data de Decolagem'].astype(str) + ' ' + df['Hora de Decolagem'].astype(str), format='mixed', dayfirst=True, errors='coerce')
    df['Data Apenas'] = df['Dep TimeStamp'].dt.date
    df['Combustível'] = pd.to_numeric(df['Combustível'], errors='coerce').fillna(0)
    
    # Merge Coordinates
    df = df.merge(df_airports[['lat', 'lon']], left_on='Localidade Decolagem', right_index=True, how='left')
    df.rename(columns={'lat': 'lat_dep', 'lon': 'lon_dep'}, inplace=True)
    df = df.merge(df_airports[['lat', 'lon']], left_on='Localidade Pouso', right_index=True, how='left')
    df.rename(columns={'lat': 'lat_arr', 'lon': 'lon_arr'}, inplace=True)
    
    # Distance & Time Computations
    def calc_dist(row):
        if pd.isna(row['lat_dep']) or pd.isna(row['lat_arr']): return 0.0
        return vincenty((row['lat_dep'], row['lon_dep']), (row['lat_arr'], row['lon_arr']))
        
    df['Distancia'] = df.apply(calc_dist, axis=1)
    df['Tempo de Voo'] = pd.to_timedelta(df['Tempo de Voo'].astype(str), errors='coerce')
    df['Tempo de Voo (h)'] = df['Tempo de Voo'].dt.total_seconds() / 3600.0

    # Build Fleet Specifications
    df_fleet = df.groupby(['Modelo', 'Emissor']).agg({
        'Combustível': 'sum',
        'Distancia': 'sum',
        'Tempo de Voo (h)': 'sum',
        'Matrícula': 'nunique'
    })
    df_fleet.rename(columns={'Matrícula': 'Tamanho da Frota'}, inplace=True)
    
    df_fleet['Consumo (km/l)'] = (df_fleet['Distancia'] / df_fleet['Combustível']).replace([float('inf'), -float('inf')], 0).fillna(0)
    df_fleet['Velocidade Média (km/h)'] = (df_fleet['Distancia'] / df_fleet['Tempo de Voo (h)']).replace([float('inf'), -float('inf')], 0).fillna(0)
    
    models = df_fleet.index.get_level_values('Modelo')
    df_fleet['Capacidade'] = models.map(lambda m: AIRCRAFT_CONFIG[m]['cap'])
    df_fleet['Alcance'] = models.map(lambda m: AIRCRAFT_CONFIG[m]['range'])
    df_fleet['Diaria_Litros'] = models.map(lambda m: AIRCRAFT_CONFIG[m]['daily_fuel'])

    # Build Demands
    df_valid = df.dropna(subset=['Data Apenas']).copy()
    df_valid = df_valid[(df_valid['Localidade Decolagem'] != '0') & (df_valid['Localidade Pouso'] != '0')]
    demands_grouped = df_valid.groupby(['Data Apenas', 'Localidade Decolagem', 'Localidade Pouso'])['PAX'].sum().reset_index()
    demands = demands_grouped[demands_grouped['PAX'] > 0].copy()

    # Load External Lookups
    with open(JSON_EMISSORES, 'r') as f:
        emissores_data = json.load(f)[0]
    
    with open(JSON_AEROPORTOS, 'r') as f:
        airports_geo = json.load(f)
        
    # Validation Lists
    valid_icaos = df_airports.index.intersection(list(set(demands['Localidade Decolagem'].unique().tolist() + demands['Localidade Pouso'].unique().tolist() + list(emissores_data.values())))).tolist()
    
    return demands, df_fleet.reset_index(), valid_icaos, airports_geo