##
# .bin.pg_python - Python console with a connection.
##
"""
Python command with a PG-API connection(``db``).
"""
import os
import sys
import re
import code
import optparse
import contextlib
from .. import clientparameters
from ..python import command as pycmd
from .. import __version__

from ..driver import default as pg_driver
from .. import exceptions as pg_exc
from .. import sys as pg_sys
from .. import lib as pg_lib

pq_trace = optparse.make_option(
	'--pq-trace',
	dest = 'pq_trace',
	help = 'trace PQ protocol transmissions',
	default = None,
)
default_options = [
	pq_trace,
	clientparameters.option_lib,
	clientparameters.option_libpath,
] + pycmd.default_optparse_options

def command(argv = sys.argv):
	p = clientparameters.DefaultParser(
		"%prog [connection options] [script] ...",
		version = __version__,
		option_list = default_options
	)
	p.disable_interspersed_args()
	co, ca = p.parse_args(argv[1:])
	rv = 1

	# Resolve the category.
	pg_sys.libpath.insert(0, os.path.curdir)
	pg_sys.libpath.extend(co.libpath or [])
	if co.lib:
		cat = pg_lib.Category(*map(pg_lib.load, co.lib))
	else:
		cat = None

	trace_file = None
	if co.pq_trace is not None:
		trace_file = open(co.pq_trace, 'a')

	try:
		need_prompt = False
		cond = None
		connector = None
		connection = None
		while connection is None:
			try:
				cond = clientparameters.collect(parsed_options = co, prompt_title = None)
				if need_prompt:
					# authspec error thrown last time, so force prompt.
					cond['prompt_password'] = True
				try:
					clientparameters.resolve_password(cond, prompt_title = 'pg_python')
				except EOFError:
					raise SystemExit(1)
				connector = pg_driver.fit(category = cat, **cond)
				connection = connector()
				if trace_file is not None:
					connection.tracer = trace_file.write
				connection.connect()
			except pg_exc.ClientCannotConnectError as err:
				for att in connection.failures:
					exc = att.error
					if isinstance(exc, pg_exc.AuthenticationSpecificationError):
						sys.stderr.write(os.linesep + exc.message + (os.linesep*2))
						# keep prompting the user
						need_prompt = True
						connection = None
						break
				else:
					# no invalid password failures..
					raise

		pythonexec = pycmd.Execution(ca,
			context = getattr(co, 'python_context', None),
			loader = getattr(co, 'python_main', None),
		)

		builtin_overload = {
		# New built-ins
			'connector' : connector,
			'db' : connection,
			'do' : connection.do,
			'prepare' : connection.prepare,

			'sqlexec' : connection.execute,
			'settings' : connection.settings,
			'proc' : connection.proc,
			'xact' : connection.xact,
		}
		if not isinstance(__builtins__, dict):
			builtins_d = __builtins__.__dict__
		else:
			builtins_d = __builtins__
		restore = {k : builtins_d.get(k) for k in builtin_overload}

		builtins_d.update(builtin_overload)
		try:
			with connection:
				rv = pythonexec(
					context = pycmd.postmortem(os.environ.get('PYTHON_POSTMORTEM'))
				)
				exc = getattr(sys, 'last_type', None)
				if rv and exc and not issubclass(exc, Exception):
					# Don't try to close it if wasn't an Exception.
					del connection.pq.socket
		finally:
			# restore __builtins__
			builtins_d.update(restore)
			for k, v in builtin_overload.items():
				if v is None:
					del builtins_d[x]
			if trace_file is not None:
				trace_file.close()
	except:
		pg_sys.libpath.remove(os.path.curdir)
		raise
	return rv

if __name__ == '__main__':
	sys.exit(command(sys.argv))
##
# vim: ts=3:sw=3:noet:
