import sqlite3
import pandas as pd
from argparse import ArgumentParser
from tqdm.rich import tqdm
from loguru import logger

def short_id(id: str):
    new_id = id.removeprefix('detail_')
    new_id = new_id.removeprefix('corner_')
    new_id = new_id.removeprefix('number_')
    return new_id

def main():
    p = ArgumentParser()
    p.add_argument('--path', type=str, help='path to pacenotes.csv')
    p.add_argument('--db', type=str, help='path to db')
    args = vars(p.parse_args())

    df = pd.read_csv(args['path'])  # pacenotes.csv
    df.sort_values(by='Id', inplace=True)
    df.reset_index(drop=True, inplace=True)

    data_pacenote = []
    data_pacenote_filenames_primary = []
    data_pacenote_filenames = []
    rows = list(df.iterrows())
    for i, row in tqdm(rows):
        data_pacenote.append((i, row['Description']))
        data_pacenote_filenames_primary.append((i, short_id(row['Id']), True))
        data_pacenote_filenames.append((i, row['Id'], False))

    logger.info(f'got {len(data_pacenote)} pacenotes and {len(data_pacenote_filenames)} filenames')    

    conn = sqlite3.connect(args['db'])
    c = conn.cursor()
    c.executemany('INSERT INTO pacenote(`id`, `description`) VALUES (?, ?)', data_pacenote)
    c.executemany('INSERT INTO pacenote_filenames(`id`, `filename`, `is_primary`) VALUES (?, ?, ?)', data_pacenote_filenames_primary)
    c.executemany('INSERT INTO pacenote_filenames(`id`, `filename`, `is_primary`) VALUES (?, ?, ?)', data_pacenote_filenames)
    conn.commit()

    logger.info('done')

if __name__ == '__main__':
    main()