import sqlite3
import json
from tqdm import tqdm
import argparse

def main():
    # Argument parsing
    parser = argparse.ArgumentParser(description='Process and insert JSON data into SQLite database.')
    parser.add_argument('db_name', type=str, help='Name of the SQLite database file.')
    parser.add_argument('input_file', type=str, help='Name of the input JSON file.')
    parser.add_argument('table_name', type=str, help='Name of the table in the database.')
    args = parser.parse_args()

    # Connect to the SQLite database
    conn = sqlite3.connect(args.db_name)
    cursor = conn.cursor()

    # Process and insert JSON data
    with open(args.input_file, 'r', encoding='utf-8') as f:
        for line in tqdm(f, desc='Processing data', unit=' lines'):
            data = json.loads(line)
            fen = data['fen']
            best_eval = max(data['evals'], key=lambda x: x['depth'])
            depth = best_eval['depth']

            cur_player = fen.split(' ')[1]

            # Safely get evaluation details
            pvs = best_eval.get('pvs', [{}])[0]
            cpe = pvs.get('cp', 20000 if cur_player == 'w' else -20000)
            mate = pvs.get('mate')
            nxt = pvs.get('line', '').split(' ')[0]

            # Insert data into the table
            cursor.execute(f'''
                INSERT OR IGNORE INTO {args.table_name} (fen, cpe, dep, nxt, mate) 
                VALUES (?, ?, ?, ?, ?)
            ''', (fen, cpe, depth, nxt, mate))

    # Commit the transaction and close the connection
    conn.commit()
    conn.close()

    print("Data processing and insertion complete.")

if __name__ == "__main__":
    main()