'''Diego Garay-Ruiz, 2022.
Processing of the CML files generated by ioChem-BD to generate Python dictionaries,
based on XSLT stylesheets.
Default stylesheets for some programs are provided in ../stylesheets, but custom ones
can also be provided.'''

import lxml.etree as ET
import os.path
from saxonche import PySaxonProcessor
import json
from importlib_resources import files

def process_xslt_entry(xslt_string,main_separator="#;#"):
	'''Process a single cc:job entry in a CML file, transforming it to a plain text string
	through XSLT and then generates a Python dictionary with all parsed fields, as strings.
	Input:
	- xslt_string. String generated by the XSLT transformation.
	- main_separator. String used in the XSL files to separate different key:value fields. Must match
	the template: provided templates will use the default "#;#" separator
	Output:
	- parsed_dict. Dict containing key:value pairs for all the fields requested by the XSLT,
	with all values being strings.'''
	raw_entries = xslt_string.split(main_separator)
	gen_dict_entries = [entry.split(":") for entry in raw_entries[:-1]]
	gen_dict_clean = [(item[0].strip(),item[1].strip()) for item in gen_dict_entries
					  if item[1].strip()]
	parsed_dict = dict(gen_dict_clean)
	return parsed_dict

def process_xslt_output(batch_xslt_string,job_separator="FETCH_JOB_END\n"):
	'''Direct processing of the output of the XSLT transformation of a CML document, possibly
	containing several cc:job entries.
	Input:
	- batch_xslt_string. String generated by the XSLT transformation.
	- job_separator. String used in the XSL file to separate the output from different cc:job entries.
	Must match the template: provided templates will use the default "FETCH_JOB_END\n" separator.
	Output:
	- proc_job_entries. List of dicts as generated by process_xslt_eintry for each cc:job'''
	# another separator: GENERAL_END
	if "GENERAL_END\n" in batch_xslt_string:
		general_entry,jobs_entries = batch_xslt_string.split("GENERAL_END\n")
		proc_job_entries = {"general":process_xslt_entry(general_entry)}
	else:
		proc_job_entries = {}
		jobs_entries = batch_xslt_string
	job_entries = jobs_entries.split(job_separator)
	if len(job_entries) > 1:
		job_entries[:] = job_entries[:-1]
	for idx,job_entry in enumerate(job_entries):
		proc_job_entries["job%d" % (idx+1)] = process_xslt_entry(job_entry)
	return proc_job_entries

def xslt_parsing(cml_file,custom_template=False,xslt_template="CML_Gaussian.xsl"):
	'''Direct parsing of CML files via XSLT stylesheets. By default resorts to the ../stylesheets
	folder containing default templates, but a custom XSL can also be passed.
	Input:
	- cml_file. String, name of the CML file to be parsed.
	- custom_template. Boolean, if True do not check the default directory but the path to the requested file,
	else consider the path parent to the module.
	- xslt_template. String, path to the XSL stylesheet.
	Output:
	- job_cml_fields. List of dicts as generated by process_xslt_entry for each cc:job, containing key:value
	pairs for all the fields requested by the XSLT, with all values being strings.'''
	if (not custom_template):
		xslt_path = str(files("py_iochem.stylesheets").joinpath(xslt_template))
	else:
		xslt_path = xslt_template
	doc = ET.parse(cml_file)
	transform = ET.XSLT(ET.parse(xslt_path))
	doc_transf = transform(doc)
	string_output = str(doc_transf)
	job_cml_fields = process_xslt_output(string_output)
	return job_cml_fields

def xslt_parsing_saxon(cml_file,custom_template=False,xslt_template="CML_Gaussian.xsl"):
	'''Direct parsing of CML files via XSLT stylesheets. By default resorts to the ../stylesheets
	folder containing default templates, but a custom XSL can also be passed.
	Input:
	- cml_file. String, name of the CML file to be parsed.
	- custom_template. Boolean, if True do not check the default directory but the path to the requested file,
	else consider the path parent to the module.
	- xslt_template. String, path to the XSL stylesheet.
	Output:
	- job_cml_fields. List of dicts as generated by process_xslt_entry for each cc:job, containing key:value
	pairs for all the fields requested by the XSLT, with all values being strings.'''
	if (not custom_template):
		xslt_path = str(files("py_iochem.stylesheets").joinpath(xslt_template))
	else:
		xslt_path = xslt_template

	with open(cml_file,"r") as finp:
		dump = finp.read()

	with PySaxonProcessor(license=False) as proc:
		doc = proc.parse_xml(xml_text=dump)
		transform = proc.new_xslt30_processor()
		doc_transf = transform.compile_stylesheet(stylesheet_file=xslt_path)
		string_output = doc_transf.transform_to_string(xdm_node=doc)
	job_cml_fields = process_xslt_output(string_output)
	return job_cml_fields

def dict_to_json(job_fields,outfile="cmldump.json"):
	with open(outfile,"w") as fout:
		json.dump(job_fields,fout)

	return None
