import sqlite3
import argparse
from typing import Literal

KOSPI = "KOSPI"
KOSDAQ = "KOSDAQ"

def queryAllCorpGroup(db: sqlite3.Connection, table: Literal["KOSPI", "KOSDAQ"]) -> list[tuple[str, str]]:
    assert table in [KOSPI, KOSDAQ]
    cursor = db.execute("SELECT Code, Name from "+ table)
    return [ (record[0],record[1]) for record in cursor]


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument("--create", action="store_true", help="create db schema")
    parser.add_argument("--read", action="store", help="read csv and commit in db")
    parser.add_argument("--clean", action="store_true", help="clean db")
    parser.add_argument("--table-name", required=True, choices=["KOSPI", "KOSDAQ"])
    args = parser.parse_args()
    if args.create:
        db = sqlite3.connect("./stock.db")
        db.execute(f"""
        CREATE TABLE IF NOT EXISTS "{args.table_name}" (
            "Code" TEXT,
            "Name" TEXT,
            PRIMARY KEY("Code"),
            FOREIGN KEY("Code") REFERENCES "KRXCorp"("Code")
        )
        """)
        db.close()
    elif args.read:
        path = args.read
        import csv
        with open(path, "r", encoding='UTF-8') as fp:
            reader = csv.reader(fp)
            #next(reader) # skip header
            # header not exist
            # collect data
            data = [ (row[0], row[1]) for row in reader]
        
        codeDict = { code:name for code, name in data}
        db = sqlite3.connect("./stock.db")
        dbData = {code:name for code, name in queryAllCorpGroup(db, args.table_name)}
        # remove common code
        commonCode = set(dbData.keys()) & set(codeDict.keys())
        for code in commonCode:
            del codeDict[code]
            del dbData[code]
        db.execute("BEGIN")
        # insert new code
        db.executemany(f"INSERT INTO {args.table_name} (Code, Name) VALUES (?, ?)", codeDict.items())
        # delete old code
        db.executemany(f"DELETE FROM {args.table_name} WHERE code = ?", dbData.keys())
        db.execute("COMMIT")
        db.commit()
        db.close()