import sqlite3
import os

DB_PATH = r'd:\Documents\trae_projects\yrtv\database\L3\L3_Features.sqlite'

def update_schema():
    if not os.path.exists(DB_PATH):
        print("L3 DB not found.")
        return

    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # Get existing columns
    cursor.execute("PRAGMA table_info(dm_player_features)")
    existing_cols = {row[1] for row in cursor.fetchall()}
    
    # List of columns to ensure exist
    # Copied from schema.sql
    required_columns = [
        # Basic
        'basic_avg_rating', 'basic_avg_kd', 'basic_avg_adr', 'basic_avg_kast', 'basic_avg_rws',
        'basic_avg_headshot_kills', 'basic_headshot_rate', 
        'basic_avg_first_kill', 'basic_avg_first_death', 'basic_first_kill_rate', 'basic_first_death_rate',
        'basic_avg_kill_2', 'basic_avg_kill_3', 'basic_avg_kill_4', 'basic_avg_kill_5',
        'basic_avg_assisted_kill', 'basic_avg_perfect_kill', 'basic_avg_revenge_kill',
        'basic_avg_awp_kill', 'basic_avg_jump_count',
        'basic_avg_mvps', 'basic_avg_plants', 'basic_avg_defuses', 'basic_avg_flash_assists',
        
        # STA
        'sta_last_30_rating', 'sta_win_rating', 'sta_loss_rating', 'sta_rating_volatility',
        'sta_time_rating_corr', 'sta_fatigue_decay',
        
        # BAT
        'bat_kd_diff_high_elo', 'bat_kd_diff_low_elo', 'bat_avg_duel_win_rate', 'bat_avg_duel_freq',
        'bat_win_rate_close', 'bat_win_rate_mid', 'bat_win_rate_far',
        
        # HPS
        'hps_clutch_win_rate_1v1', 'hps_clutch_win_rate_1v2', 'hps_clutch_win_rate_1v3_plus',
        'hps_match_point_win_rate', 'hps_undermanned_survival_time', 'hps_pressure_entry_rate',
        'hps_momentum_multikill_rate', 'hps_tilt_rating_drop', 'hps_clutch_rating_rise',
        'hps_comeback_kd_diff', 'hps_losing_streak_kd_diff',
        
        # PTL
        'ptl_pistol_kills', 'ptl_pistol_multikills', 'ptl_pistol_win_rate', 'ptl_pistol_kd', 'ptl_pistol_util_efficiency',
        
        # SIDE
        'side_rating_ct', 'side_rating_t', 'side_kd_ct', 'side_kd_t',
        'side_win_rate_ct', 'side_win_rate_t',
        'side_first_kill_rate_ct', 'side_first_kill_rate_t',
        'side_kd_diff_ct_t',
        'side_kast_ct', 'side_kast_t',
        'side_rws_ct', 'side_rws_t',
        'side_first_death_rate_ct', 'side_first_death_rate_t',
        'side_multikill_rate_ct', 'side_multikill_rate_t',
        'side_headshot_rate_ct', 'side_headshot_rate_t',
        'side_defuses_ct', 'side_plants_t',
        'side_obj_ct', 'side_obj_t',
        'side_planted_bomb_count', 'side_defused_bomb_count',
        
        # UTIL
        'util_avg_nade_dmg', 'util_avg_flash_time', 'util_avg_flash_enemy', 'util_avg_flash_team', 'util_usage_rate',
        
        # Scores
        'score_bat', 'score_sta', 'score_hps', 'score_ptl', 'score_tct', 'score_util'
    ]
    
    for col in required_columns:
        if col not in existing_cols:
            print(f"Adding missing column: {col}")
            try:
                # Most are REAL, integers are fine as REAL in sqlite usually, or use affinity
                cursor.execute(f"ALTER TABLE dm_player_features ADD COLUMN {col} REAL")
            except Exception as e:
                print(f"Failed to add {col}: {e}")
                
    conn.commit()
    conn.close()
    print("Schema update check complete.")

if __name__ == "__main__":
    update_schema()