"""
Helper functions and classes for unit tests and similar.

Whatever is useful to unit tests from here should be imported into
testhelpers, too.  Unit test modules should not be forced to import
this.
"""

#c Copyright 2008-2021, the GAVO project <gavo@ari.uni-heidelberg.de>
#c
#c This program is free software, covered by the GNU GPL.  See the
#c COPYING file in the source distribution.


import contextlib
import gzip
import io
import os
import re
import tempfile

from lxml import etree

from twisted.internet  import defer
from twisted.python import failure
from twisted.python import urlpath
from twisted.web import resource
from twisted.web import server
from twisted.web.test.requesthelper import DummyRequest

from gavo import base
from gavo import utils
from gavo import web
from gavo.utils import stanxml


def _nukeNamespaces(xmlString):
	"""removes namespace declarations from xmlString (which must be bytes).

	This is for writing more compact tests and should of course not be
	used outside of tests; in particular, you could easily fool the
	mechanism to wreck your XML.

	This always returns bytes.
	"""
	nsCleaner = re.compile(rb'^(</?)(?:[a-z0-9]+:)')
	return re.sub(b"(?s)<[^>]*>", 
		lambda mat: nsCleaner.sub(rb"\1", mat.group()),
		re.sub(b'xmlns="[^"]*"', b"", xmlString))


class _WrappedEtree:
	"""a wrapper adding a few methods to an lxml etree.

	This is done as a wrapper because you can't monkeypatch _Element.

	See getXMLTree for what this about; it's essentially an implemenation
	detail of that function.
	"""
	def __init__(self, tree):
		self._tree = tree
	
	def __getattr__(self, name):
		return getattr(self._tree, name)

	def __getitem__(self, index):
		return self._tree[index]

	def uniqueXpath(self, path):
		res = self.xpath(path)
		assert len(res)==1, "Xpath %s gave %d matches"%(path, len(res))
		return res[0]

	def getById(self, id):
		return self.uniqueXpath("//*[@id='%s']"%id)
	
	def getByID(self, id):
		return self.uniqueXpath("//*[@ID='%s']"%id)


def getXMLTree(xmlString, debug=False):
	"""returns an ``libxml2`` etree-like object for ``xmlString``, where, 
	for convenience, all namespaces on elements are nuked.

	This will only accept strings.

	The libxml2 etree lets you do xpath searching using the ``xpath`` method.

	Nuking namespaces is of course not a good idea in general, so you
	might want to think again before you use this in production code.

	To facilitate writing tests, in addition to lxml.etree methods the returned
	object also has the following methods:

	* uniqueXpath(xpath), returning a single match if it's there
	  and raises an assertion error otherwise.
	* getById(id), returning the unique element with id and raising an 
	  assertion error if that doesn't exist.
	* getByID(id), as getById, by for VOTable-style ID.
	"""
	tree = etree.fromstring(_nukeNamespaces(xmlString))

	if debug:
		etree.dump(tree)
	return _WrappedEtree(tree)


class XSDResolver(etree.Resolver):
	"""A resolver for external entities only returning in-tree files.
	"""
	def __init__(self):
		self.basePath = "schemata"

	def getPathForName(self, name):
		xsdName = name.split("/")[-1]
		return base.getPathForDistFile(
			os.path.join(self.basePath, xsdName))

	def resolve(self, url, pubid, context):
		try:
			# resolve namespace URIs, too
			try:
				url = stanxml.NSRegistry.getSchemaForNS(url)
			except base.NotFoundError:
				# it's not a (known) namespace URI, try on
				pass
			
			path = self.getPathForName(url)
			res = self.resolve_filename(path, context)
			if res is not None:
				return res
		except:
			pass # fall through to error message
		base.ui.notifyError("Did not find local file for schema %s --"
			" this will fall back to network resources and thus probably"
			" be slow"%url)


RESOLVER = XSDResolver()
XSD_PARSER = etree.XMLParser()
XSD_PARSER.resolvers.add(RESOLVER)


@contextlib.contextmanager
def MyParser():
	if etree.get_default_parser is XSD_PARSER:
		yield
	else:
		etree.set_default_parser(XSD_PARSER)
		try:
			yield
		finally:
			etree.set_default_parser()

class QNamer(object):
	"""A hack that generates QNames through getattr.

	Construct with the desired namespace.
	"""
	def __init__(self, ns):
		self.ns = ns
	
	def __getattr__(self, name):
		return etree.QName(self.ns, name.strip("_"))

XS = QNamer("http://www.w3.org/2001/XMLSchema")


VO_SCHEMATA = [
		"b2find.xsd",
		"Characterisation.xsd",
		"Colstats.xsd",    # remove once it's in VODataService
		"ConeSearch.xsd",
		"DaCHS.xsd",
		"DataModel.xsd",
		"DocRegExt.xsd",
		"oai_dc.xsd",
		"OAI-PMH.xsd",
		"RegistryInterface.xsd",
		"SIA.xsd",
		"SLAP.xsd",
		"SSA.xsd",
		"StandardsRegExt.xsd",
		"stc.xsd",
		"stc-v1.20.xsd",
		"coords-v1.20.xsd",
		"region-v1.20.xsd",
		"TAPRegExt.xsd",
		"UWS.xsd",
		"VODataService.xsd",
		"VOEvent.xsd",
		"VOEventRegExt.xsd",
		"VORegistry.xsd",
		"VOResource.xsd",
		"VOSIAvailability.xsd",
		"VOSICapabilities.xsd",
		"VOSITables.xsd",
		"VOTable-1.1.xsd",
		"VOTable-1.2.xsd",
		"VOTable.xsd",
		"vo-dml.xsd",
		"xlink.xsd",
		"XMLSchema.xsd",
		"xml.xsd",]


def getJointValidator(schemaPaths):
	"""returns an lxml validator containing the schemas in schemaPaths.

	schemaPaths must be actual file paths, absolute or
	trunk/schema-relative.
	"""
	with MyParser():
		subordinates = []
		for fName in schemaPaths:
			fPath = RESOLVER.getPathForName(fName)
			root = etree.parse(fPath).getroot()
			subordinates.append((
				"http://vo.ari.uni-heidelberg.de/docs/schemata/"+fName,
				root.get("targetNamespace")))

		root = etree.Element(
			XS.schema, attrib={"targetNamespace": "urn:combiner"})
		for schemaLocation, tns in subordinates:
			etree.SubElement(root, XS.import_, attrib={
				"namespace": tns, "schemaLocation": schemaLocation})
		
		doc = etree.ElementTree(root)
		return etree.XMLSchema(doc)


def getDefaultValidator(extraSchemata=[]):
	"""returns a validator that knows the schemata typically useful within
	the VO.

	This will currently only work if DaCHS is installed from an SVN
	checkout with setup.py develop.

	What's returned has a method assertValid(et) that raises an exception 
	if the elementtree et is not valid.  You can simply call it to
	get back True for valid and False for invalid.
	"""
	return getJointValidator(VO_SCHEMATA+extraSchemata)


def _makeLXMLValidator():
	"""returns an lxml-based schema validating function for the VO XSDs

	This is not happening at import time as it is time-consuming, and the 
	DaCHS server probably doesn't even validate anything.

	This is used below to build getXSDErrorsLXML.
	"""
	VALIDATOR = getDefaultValidator()

	def getErrors(data, leaveOffending=False):
		"""returns error messages for the XSD validation of the string in data.
		"""
		try:
			with MyParser():
				if hasattr(data, "xpath"):
					# we believe it's already parsed stuff
					tree = data
				else:
					tree = etree.fromstring(data)

				if VALIDATOR.validate(tree):
					return None
				else:
					if leaveOffending:
						if hasattr(data, "xpath"):
							data = etree.tostring(data, encoding="utf-8")
						with open("badDocument.xml", "wb") as of:
							of.write(data)
					return str(VALIDATOR.error_log)
		except Exception as msg:
			return str(msg)
	
	return getErrors


def getXSDErrorsLXML(data, leaveOffending=False):
	"""returns error messages for the XSD validation of the string in data.
	"""
	if not hasattr(getXSDErrorsLXML, "validate"):
		getXSDErrorsLXML.validate = _makeLXMLValidator()
	return getXSDErrorsLXML.validate(data, leaveOffending)


getXSDErrors = getXSDErrorsLXML


class XSDTestMixin(object):
	"""provides a assertValidates method doing XSD validation.

	assertValidates raises an assertion error with the validator's
	messages on an error.  You can optionally pass a leaveOffending
	argument to make the method store the offending document in
	badDocument.xml.
	"""
	def assertValidates(self, xmlSource, leaveOffending=False):
		messages = getXSDErrors(xmlSource, leaveOffending)
		if messages:
			raise AssertionError(messages)

	def assertWellformed(self, xmlSource):
		try:
			etree.fromstring(xmlSource)
		except Exception as msg:
			raise AssertionError("XML not well-formed (%s)"%msg)


def getMemDiffer(ofClass=base.Structure):
	"""returns a function to call that returns a list of new DaCHS structures
	since this was called.

	If you watch everything, things get hairy because of course the state
	of this function (for instance) also creates references.  Hence, pass
	ofClass to choose what the funtion will track.

	This will call a gc.collect itself (and wouldn't make sense without that)
	"""
	import gc
	gc.collect()
	seen_ids = set()
	for ob in gc.get_objects():
		try:
			if isinstance(ob, ofClass):
				seen_ids.add(id(ob))
		except ReferenceError:
			# object is already essentially gone, don't worry about it.
			pass
	del ob
	
	def getNewObjects():
		gc.collect()
		newObjects = []
		for ob in gc.get_objects():
			try:
				if id(ob) not in seen_ids and isinstance(ob, ofClass):
					newObjects.append(ob)
			except ReferenceError:
				# again, don't worry about disappearing objects
				pass

		return newObjects
	
	return getNewObjects


def getUnreferenced(items):
	"""returns a list of elements in items that do not have a reference
	from any other in items.
	"""
	import gc
	itemids = set(id(i) for i in items)
	unreferenced = []
	for i in items:
		intrefs = set(id(r) for r in gc.get_referrers(i)) & itemids
		if not intrefs:
			unreferenced.append(i)
	return unreferenced

		
def debugReferenceChain(ob):
	"""a sort-of-interactive way to investigate where ob is referenced.

	* d -- enter pdb (look at ob, perhaps at nob)
	* u -- follow 
	* x -- continue execution
	"""
	import gc
	while True:
		print("Current object: ", repr(ob))
		refs = gc.get_referrers(ob)
		if not refs:
			print("Not referenced -- exiting")
			break

		while refs:
			nob = refs.pop()
			print(len(refs), utils.makeEllipsis(repr(nob)))
			res = input("?")
			
			if res=="d":
				import pdb;pdb.Pdb(nosigint=True).set_trace()

			elif res=="x":
				return

			elif res=="u":
				ob = nob
				break

			elif res=="?":
				print("d, x, u, <empty>")

			elif not refs:
				print("Referrers exhausted, warping")


NEWIDS = set()

def memdebug(watchClass=base.Structure):
	"""a debug method to track memory usage after some code has run.

	This is typically run from ArchiveService.getChild, since request
	processing should be idempotent wrt memory after initial caching.

	This is for editing in place by DaCHS plumbers; accordingly, you're
	not supposed to make sense of this.
	"""
	import gc
	print(">>>>>> total managed:", len(gc.get_objects()))

	if hasattr(base, "getNewStructs"):
		ns = base.getNewStructs()
		print(">>>>>> new objects:", len(ns))
		if len(ns)<11000:
			ur = getUnreferenced(ns)
			print(">>>>>> new externally referenced:", len(ur))
			del ur
			print([ob for ob in ns if isinstance(ob, watchClass)])
			if True:
				try:
					debugReferenceChain(
						[ob for ob in ns if isinstance(ob, watchClass)][0])
				except IndexError:
					pass

	base.getNewStructs = getMemDiffer(ofClass=watchClass)


@contextlib.contextmanager
def testFile(name, 
		content, 
		writeGz=False, 
		inDir=base.getConfig("tempDir"),
		timestamp=None):
	"""a context manager that creates a file name with content in inDir.

	The full path name is returned.

	content can be bytes or str; in the latter case, it's utf-8 encoded
	before writing.

	With writeGz=True, content is gzipped on the fly (don't do this if
	the data already is gzipped).

	You can pass in name=None to get a temporary file name if you don't care
	about the name.

	inDir will be created as a side effect if it doesn't exist but (right
	now, at least), not be removed.
	"""
	if not os.path.isdir(inDir):
		os.makedirs(inDir)

	if name is None:
		handle, destName = tempfile.mkstemp(dir=inDir)
		os.close(handle)
	else:
		destName = os.path.join(inDir, name)

	if writeGz:
		f = gzip.GzipFile(destName, mode="wb")
	else:
		f = open(destName, "wb")

	f.write(utils.bytify(content))
	f.close()
	
	if timestamp:
		os.utime(destName, times=(timestamp, timestamp))

	try:
		yield destName
	finally:
		try:
			os.unlink(destName)
		except os.error:
			pass



@contextlib.contextmanager
def collectedEvents(*kinds):
	"""a context manager collecting event arguments for a while.

	The yielded thing is a list that contains tuples of event name and
	the event arguments.
	"""
	collected = []

	def makeHandler(evType):
		def handler(*args):
			collected.append((evType,)+args)
		return handler
	
	handlers = [(kind, makeHandler(kind)) for kind in kinds]
	for kind, handler in handlers:
		base.ui.subscribe(kind, handler)

	try:
		yield collected
	finally:
		for kind, handler in handlers:
			base.ui.unsubscribe(kind, handler)
		

class FakeFile:
	def __init__(self, filename, fileobj):
		self.filename, self.file = filename, fileobj
		if isinstance(self.file, bytes):
			self.file = io.BytesIO(self.file)


class FakeFieldStorage(object):
	filename = None
	def __init__(self, args):
		self.args = {}
		for k, vals in args.items():
			res = self.args[k.decode("utf-8")] = []
			for v in vals:
				if isinstance(v, bytes):
					res.append(v.decode("utf-8"))
				else:
					res.append(FakeFile(*v))

	def __iter__(self):
		return iter(self.args)
	
	def getfirst(self, key):
		return self.args[key][0]
	
	def getlist(self, key):
		return self.args[key]

	def __getitem__(self, key):
		return self.args[key][0]

	def keys(self):
		return list(self.args.keys())

	@property
	def file(self):
		return io.StringIO(self.args[0])


def bytify_seq(s):
	"""returns a list from s with any elements bytified.

	s must be a sequence or None.  Oh, for convenience we also accept plain
	bytes and strings.
	"""
	if s is None:
		return None
	if isinstance(s, (bytes, str)):
		s = [s]
	if not isinstance(s, (list, tuple)):
		raise Exception("bytify_seq really wants a sequence, not %s"%repr(s))
	return [utils.bytify(item) for item in s]


class FakeRequest(DummyRequest, web.Request):
	"""A request for test purposes.

	We furnish t.w's DummyRequest with some extra facilities to let us
	be a bit lazy in having rather macro tests.

	Also, stock twisted DummyRequest produces an endless loop with push
	producers (which is what we have), so we fix that, too.
	"""

	method = b"GET"
	session = None

	# some code tests for a live connection using client
	client = True

	def __init__(self, uri=b'', headers=None, args=None, avatar=None,
                 currentSegments=None, cookies=None,
                 user="", password="", isSecure=False):
		uri = utils.bytify(uri)
		if uri.startswith(b"/"):
			uri = uri[1:]
		postpath = []
		if uri:
			postpath = uri.split(b"/")

		DummyRequest.__init__(self, uri)

		if args:
			self.args = dict((utils.bytify(k), bytify_seq(v)) 
				for k, v in args.items())
		else:
			self.args = {}

		self.uri = uri
		self.postpath = postpath
		self.code = 200
		self.user, self.password = user, password
		self.deferred = defer.Deferred()
		self.accumulator = b""
		self.prepath = []
		self.fields = FakeFieldStorage(self.args)
		self.finished = False
		self.secure = False
		self.channel = 1 # must be non-None for custom hangup detection

	def setHeader(self, name, value):
# overridden because t.w.t doesn't overwrite
		self.responseHeaders.setRawHeaders(name, [value])

	def write(self, data):
		self.accumulator += utils.bytify(data)

	def notifyFinish(self):
		return self.deferred

	def prePathURL(self):
		return 'http://%s/%s'%(self.getHost(), '/'.join(self.prepath))

	def finish(self):
		self.finished = True
		self.deferred.callback((self.accumulator, self))

	def finishCallback(self, arg):
		if isinstance(arg, failure.Failure):
			arg.raiseException()
		if not self.finished:
			self.finish()

	def setHost(self, host, port):
		self.host = host
	
	def getHost(self):
		return self.host

	def setResponseCode(self, code):
		self.code = code

	def URLPath(self):
		return urlpath.URLPath.fromString(self.path.decode("utf-8"))

	@property
	def path(self):
		return self.uri

	def isSecure(self):
		return self.secure

	def getLocationValue(self):
		"""returns a location header if this requests redirects, and raises
		an AssertionError otherwise.
		"""
		if not self.code or self.code//100!=3:
			raise AssertionError("Trying to get a redirection target for"
				" request with status %s"%self.code)

		hdrs = self.responseHeaders.getRawHeaders("location")
		if len(hdrs)!=1:
			raise AssertionError("Trying to get a redirection target when"
				" location headers are %s"%hdrs)

		return hdrs[0]

	def processWithRoot(self, page):
		"""runs this request on page.

		This is probably a bad idea all around, and we should just be using
		trial.  But since sync tests are quite a bit more convenient,
		here this is.  Of course, it only works if resource effectively
		renders sync (or has a renderSync method).
		"""
		rsc = resource.getChildForRequest(page, self)

		res = getattr(rsc, "renderSync", rsc.render)(self)
		if res:
			if isinstance(res, int) and res==server.NOT_DONE_YET:
				# this will only work if the thing is actually sync.
				# see servicetest._syncvosi for an inspration there.
				# But in that case, accumulator will have it all.
				pass
			else:
				return res

		return self.accumulator

	def registerProducer(self, producer, isPush):
		self.producer = producer
		if not isPush:
			DummyRequest.registerProducer(
				self, producer, isPush)

	def unregisterProducer(self):
		# stop twisted pull producers, too
		self.go = 0
		self.channel = None
		del self.producer
	
	def getResponseHeader(self, headerName):
		return self.responseHeaders.getRawHeaders(headerName)[0]
