#   Copyright © 2023  Lee Yingtong Li (RunasSudo)
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU Affero General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU Affero General Public License for more details.
#
#   You should have received a copy of the GNU Affero General Public License
#   along with this program.  If not, see <https://www.gnu.org/licenses/>.

import pandas as pd
import sqlite3
import zipfile

con = sqlite3.connect('database.db')
cur = con.cursor()

# Init schema
cur.execute('DROP TABLE IF EXISTS pbs_drug')
cur.execute('CREATE TABLE pbs_drug (id INTEGER PRIMARY KEY AUTOINCREMENT, item_code CHARACTER(6), mp_pt TEXT, tpuu_or_mpp_pt TEXT, restriction_flag CHARACTER(1), mq INTEGER, repeats INTEGER)')

cur.execute('DROP TABLE IF EXISTS pbs_prescriber_type')
cur.execute('CREATE TABLE pbs_prescriber_type (id INTEGER PRIMARY KEY AUTOINCREMENT, item_code CHARACTER(6), prescriber_type CHARACTER(1))')

cur.execute('DROP TABLE IF EXISTS pbs_streamlined')
cur.execute('CREATE TABLE pbs_streamlined (id INTEGER PRIMARY KEY AUTOINCREMENT, item_code CHARACTER(6), treatment_of_code INTEGER)')

# Read drug list, prescriber type
with zipfile.ZipFile('data/2023-01-01-v3extracts.zip', 'r') as zipf:
	# drug_xxx.txt
	
	with zipf.open('drug_20230101.txt', 'r') as f:
		df_drug = pd.read_csv(f, sep='!')
	
	for _, drug in df_drug.iterrows():
		# Skip already added
		cur.execute('SELECT COUNT(*) FROM pbs_drug WHERE item_code=?', (drug['item-code'],))
		if cur.fetchone()[0] > 0:
			continue
		
		cur.execute('INSERT INTO pbs_drug (item_code, mp_pt, tpuu_or_mpp_pt, restriction_flag, mq, repeats) VALUES (?, ?, ?, ?, ?, ?)', (drug['item-code'], drug['mp-pt'], drug['tpuu-or-mpp-pt'], drug['restriction-flag'], drug['mq'], drug['repeats']))
	
	# Prescriber_type_xxx.txt
	
	with zipf.open('Prescriber_type_20230101.txt', 'r') as f:
		df_prescriber_type = pd.read_csv(f, sep='\t', header=0, names=['mp-pt', 'item-code', 'prescriber-type'])
	
	for _, prescriber_type in df_prescriber_type.iterrows():
		cur.execute('INSERT INTO pbs_prescriber_type (item_code, prescriber_type) VALUES (?, ?)', (prescriber_type['item-code'], prescriber_type['prescriber-type']))
	
	# streamlined_xxx.txt (streamlined authorities)
	
	with zipf.open('streamlined_20230101.txt', 'r') as f:
		df_streamlined = pd.read_csv(f, sep='\t')
	
	for _, streamlined in df_streamlined.iterrows():
		cur.execute('INSERT INTO pbs_streamlined (item_code, treatment_of_code) VALUES (?, ?)', (streamlined['item-code'], streamlined['treatment-of-code']))

con.commit()