import csv, sqlite3, json

conn = sqlite3.connect("/root/tsk.db")
c = conn.cursor()
c.execute("DROP TABLE IF EXISTS tsk_data")
c.execute("CREATE TABLE tsk_data (id INTEGER PRIMARY KEY, source TEXT, id_code TEXT, title TEXT, content TEXT)")
conn.commit()

def imp(path, src, tcol, icol, ccol):
    with open(path, encoding="utf-8-sig") as f:
        rdr = csv.DictReader(f)
        n = 0
        for row in rdr:
            t = (row.get(tcol) or "").strip()
            i = (row.get(icol) or "").strip()
            ct = (row.get(ccol) or "").strip()
            c.execute("INSERT INTO tsk_data(source,id_code,title,content) VALUES(?,?,?,?)",
                      (src, i, t, ct))
            n += 1
            if n % 2000 == 0:
                conn.commit()
        conn.commit()
        print(f"  {src}: {n}", flush=True)

imp("/root/pharma.csv", "pharma", "名称", "ID", "原文")
imp("/root/cpi.csv", "cpi", "项目名称", "ID", "项目原文")
imp("/root/plant.csv", "plant", "公司名称", "Plant ID", "原文信息")

# Create FTS
c.execute("DROP TABLE IF EXISTS tsk_fts")
c.execute("CREATE VIRTUAL TABLE tsk_fts USING fts5(title, content, id_code)")
c.execute("INSERT INTO tsk_fts(rowid, title, content, id_code) SELECT id, title, content, id_code FROM tsk_data")
conn.commit()

t = c.execute("SELECT count(*) FROM tsk_data").fetchone()[0]
print(f"\n总计: {t} 条", flush=True)
conn.close()