"""
# Add file restriction in upload area - Done
# Fix summary hanging issue
# Custom highlight function
# Slider for number of keywords - with re-running
# Validate input - check for empty input - first run vs later runs, 
#	in later runs, state variable will not be empty, so need to consider that case 
"""

import streamlit as st
from streamlit_tags import st_tags
# import annotated_text as at
from annotated_text import annotated_text

from io import StringIO
import pandas as pd

import read_extract
import re
import threading
import time
import os

def check(datatype, task, field, summ=True):
	"""
	Function that calls the necessary internal functions for summary and keyword extraction
	"""
	with result_placeholder:
		with st.spinner('Please wait while we process the keywords:'):
			if summ:
				st.session_state['summariser'] = read_extract.AbstractiveSummarizer()
				st.session_state['summ_thread'] = threading.Thread(target=st.session_state['summariser'].bart, args=(st.session_state['lic_txt'],))
				st.session_state['summ_thread'].start()

			res = read_extract.get_keywords(datatype, task, field, st.session_state['pos_text'], st.session_state['neg_text'])
			

			st.session_state['p_keywords'], st.session_state['n_keywords'], st.session_state['contrd'], st.session_state['hl_text'] = res



def display(ph, title, content, highlight=False):

	"""
	Helper function to display the different contents based on the button presses
	"""
	with ph.container():
		st.markdown('### '+title)
		# st.markdown('---')
		if highlight:
			annotated_text(*content)
			# highlight_cont(content)
		else:
			st.markdown(content)

def output():

	"""
	Function to display the final output after extracting Summary and keywords
	"""
	if st.session_state['p_keywords'] or st.session_state['n_keywords']:
		with result_placeholder.container():
			st.markdown('---')
			msg_placeholder = st.empty()
			col1, col2, col3, col4 = st.columns(4)
			content_placeholder = st.empty()

			with msg_placeholder:
				if not st.session_state.contrd:
					st.success("### Congrats!! There are no contradictions to the license.")
				else:
					st.error("### The chosen usage contradicts with the license.")
					with col4:
						st.session_state['op4'] = st.button(label='View Usage Contradictions', type='secondary')#,
															   # on_click=display,
															   # args=(content_placeholder, 'Contradictions', st.session_state['hl_text'], True))
						if st.session_state['op4']:
							display(content_placeholder, 'Contradictions', st.session_state['hl_text'], True)
			with col1:
				if st.session_state['summ']:
					st.session_state['op1'] = st.button(label='View License Summary', type='secondary')#,
														   # on_click=display,
														   # args=(content_placeholder, 'License Summary', st.session_state['summary']))
					if st.session_state['op1']:
						if (st.session_state['summ_thread']):
							if (not st.session_state['summ_thread'].is_alive()):
								with content_placeholder:
									st.markdown('The summary will take longer to generate. Please retry after sometime')
									st.session_state['op1'] = st.button(label='Retry')
							else:
								st.session_state['summ_thread'] = None
								st.session_state['summary'] = st.session_state['summariser'].summary
								display(content_placeholder, 'License Summary', st.session_state['summary'])

			with col2:
				st.session_state['op2'] = st.button(label='View Permitted Usage Tags', type='secondary')#,
													   # on_click=display,
													   # args=(content_placeholder, 'Permitted Usage Tags', ', '.join(st.session_state['p_keywords'])))
				if st.session_state['op2']:
					res = ', '.join(st.session_state['p_keywords']) if st.session_state['p_keywords'] else 'No permissions allowed'
					display(content_placeholder, 'Permitted Usage Tags', res)
			
			with col3:
				st.session_state['op3'] = st.button(label='View Restriction Tags', type='secondary')#,
													   # on_click=display,
													   # args=(content_placeholder, 'Restriction Tags', ', '.join(st.session_state['n_keywords'])))
				if st.session_state['op3']:
					res = ', '.join(st.session_state['n_keywords']) if st.session_state['n_keywords'] else 'No usage restrictions, free to use'
					display(content_placeholder, 'Restriction Tags', res)

			st.button(label='Check New Use Case', type='primary')#, on_click=reset, use_container_width=True)

def read_license(spdx):
	"""
	Function to read the processes and unprocesses license texts
	Licenses_split.csv contains the parts of the data manually separated into Dos and Dont's,
	to generate permitted(pos) and restricted(neg) use tags

	Argument:
		spdx: id used to identify the licenses. The license files are named with the spdx id for easier reading and access

	Returns / Added to session state:
		pos_text: The part of the License containing information for permitted use
		neg_text: The part of the License containing information about usage restrictions
		lic_txt: The full license text 
	"""
	try:
		file_name = spdx+'.txt'
		loc = os.path.join(os.getcwd(), 'Licenses', file_name)
		f = open(loc, encoding="utf-8")
		st.session_state.lic_txt = f.read()
	except FileNotFoundError:
		msg = 'Error: {} not found in {}'.format(file_name, os.path.dirname(loc))
		msg += """  

		  Please make sure the Licenses folder is in the same directory as app.py"""
		st.error(msg)
		reset()
	
	try:
		loc1 = os.path.join(os.getcwd(), 'Licenses_split.csv')
		loc2 = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'Licenses_split.csv')

		loc = loc1 if os.path.exists(loc1) else loc2
		sheet = pd.read_csv(loc, index_col='spdx')

		sheet = sheet.fillna('')
		st.session_state.pos_text = sheet.loc[spdx,'Dos']
		st.session_state.neg_text = sheet.loc[spdx,'Donts']

	except FileNotFoundError:
		msg = 'Error: {} not found in {}'.format('Licenses_split.csv', os.path.dirname(loc))
		msg += """  

		  Please make sure the Licenses_split.csv file is in the same directory as app.py"""
		st.error(msg)

		reset()




def get_input():

	"""
	Function to get all the necessary input for the program
	"""
	st.title("License Permissions Checker")
	st.markdown("This mini-app extracts keywords from Software License texts to check valid usage of the models")

	# Input Fields

	asset_type = st.selectbox(label='Select the type of asset being used with the License:', 
							  options=['Data', 'Model', 'Source Code', 'Model Derivatives'],
							  index=1, key="asset")
	task = st.selectbox(label='Select the task to be performed:', 
						options=['Text Classification', 'Token Classification', 'Table Question Answering', 
								'Question Answering', 'Zero-Shot Classification', 'Translation', 'Summarization', 
							  	'Conversational', 'Feature Extraction', 'Text Generation', 'Text2Text Generation', 
							  	'Fill-Mask', 'Sentence Similarity', 'Text-to-Speech', 'Automatic Speech Recognition', 
							  	'Audio-to-Audio', 'Audio Classification', 'Voice Activity Detection', 
							  	'Depth Estimation', 'Image Classification', 'Object Detection', 'Image Segmentation', 
							  	'Text-to-Image', 'Image-to-Text', 'Image-to-Image', 'Unconditional Image Generation', 
							  	'Video Classification', 'Reinforcement Learning', 'Robotics', 'Tabular Classification', 
							  	'Tabular Regression', 'Text-to-Video', 'Visual Question Answering', 'Document Question Answering', 
							  	'Zero-Shot Image Classification', 'Graph Machine Learning'],
						index=0, key="task")

	field = st_tags(label='Select the associated field (Eg. Medical, Research, Commercial, Non commercial, etc.):',
					suggestions=['Medical Conditions', 'Research', 'Commercial Use', 'Non Commercial Use', 'Criminal Likelihood Prediction', 'Synthesize Media', 'Insurance claims prediction'],
					maxtags = 1, text='', key='field')
	if field:
		field = field[0]

	lic_names = ['OpenRAIL-S', 'BigScience RAIL License v1.0', 'CreativeML OpenRAIL-M', 'BigScience BLOOM RAIL 1.0', 'Academic Free License v3.0', 'Apache license 2.0', 'Artistic license 2.0', 'Boost Software License 1.0', 'BSD 1-clause', 'BSD 2-clause "Simplified" license', 'BSD 3-clause "New" or "Revised" license', 'BSD 3-clause Clear license', 'Computational Use of Data Agreement', 'Creative Commons Zero v1.0 Universal', 'Creative Commons Attribution 4.0', 'Creative Commons Attribution Share Alike 4.0', 'Creative Commons Attribution Non Commercial 4.0', 'Creative Commons Attribution No Derivatives 4.0', 'Creative Commons Attribution Non Commercial No Derivatives 4.0', 'Creative Commons Attribution Non Commercial Share Alike 4.0', 'Educational Community License v2.0', 'Eclipse Public License 1.0', 'Eclipse Public License 2.0', 'European Union Public License 1.2', 'GNU Affero General Public License v3.0', 'GNU Free Documentation License family', 'GNU General Public License v2.0', 'GNU General Public License v3.0', 'GNU Lesser General Public License v2.1', 'GNU Lesser General Public License v3.0', 'ISC', 'LaTeX Project Public License v1.3c', 'Microsoft Public License', 'MIT', 'Mozilla Public License 2.0', 'Open Data Commons License Attribution family', 'Open Database License family', 'Open Rail++-M License', 'Open Software License 3.0', 'PostgreSQL License', 'University of Illinois/NCSA Open Source License', 'The Unlicense', 'zLib License', 'Open Data Commons Public Domain Dedication and License']
	lic_ids = ['openrail-s', 'bigscience-bloom-rail-1.0', 'creativeml-openrail-m', 'bigscience-bloom-rail-1.0', 'afl-3.0', 'apache-2.0', 'artistic-2.0', 'bsl-1.0', 'bsd-1-clause', 'bsd-2-clause', 'bsd-3-clause', 'bsd-3-clause-clear', 'c-uda', 'cc0-1.0', 'cc-by-4.0', 'cc-by-sa-4.0', 'cc-by-nc-4.0', 'cc-by-nd-4.0', 'cc-by-nc-nd-4.0', 'cc-by-nc-sa-4.0', 'ecl-2.0', 'epl-1.0', 'epl-2.0', 'eupl-1.2', 'agpl-3.0', 'gfdl', 'gpl-2.0', 'gpl-3.0', 'lgpl-2.1', 'lgpl-3.0', 'isc', 'lppl-1.3c', 'ms-pl', 'mit', 'mpl-2.0', 'odc-by', 'odbl', 'openrail++', 'osl-3.0', 'postgresql', 'ncsa', 'unlicense', 'zlib', 'pddl']

	lic_ind = st.selectbox(label='Select the License to check against:', 
							options=list(range(len(lic_ids))),
							index=0, key='license',
							format_func=lambda opt:lic_names[opt])
	spdx = lic_ids[lic_ind]

	read_license(spdx)

	summ = st.checkbox(label='Generate Summary', value=True, key='summ')

	if st.session_state.lic_txt and (st.session_state.pos_text or st.session_state.neg_text):
		submit = st.button(label='Check permissions',
						   type = 'secondary',
						   on_click=check,
							args=(asset_type, task, field, summ), )


def reset():
	"""
	Function to reset all state variables to check a new use case
	"""
	st.session_state['pos_text'] = ""
	st.session_state['neg_text'] = ""
	st.session_state['lic_txt'] = ""
	st.session_state['summary'] = ""
	st.session_state['p_keywords'] = []
	st.session_state['n_keywords'] = [] 
	st.session_state['contrd'] = False
	st.session_state['hl_text'] = ""
	for op in ['op1', 'op2', 'op3', 'op4']:
		st.session_state[op] = False
	st.session_state['summ_thread'] = None





st.set_page_config(page_title="License Usage Checker")

for key in ['pos_text', 'neg_text', 'lic_txt', 'summary', 'hl_text']:
	if key not in st.session_state:
		st.session_state[key] = ""

if "p_keywords" not in st.session_state:
	st.session_state['p_keywords'] = []

if "n_keywords" not in st.session_state:
	st.session_state['n_keywords'] = [] 

if "contrd" not in st.session_state:
	st.session_state['contrd'] = False

for op in ['op1', 'op2', 'op3', 'op4']:
	if op not in st.session_state:
		st.session_state[op] = False

if "summ_thread" not in st.session_state:
	st.session_state['summ_thread'] = None

get_input()

result_placeholder = st.empty()

output()