import logging
from trioexplorer.lib.base import *
from pylons import Response
from pylons.templating import render
import pgdb
from trio import triodb
from trio.common import TYPE_UNCERTAIN, TYPE_CONFIDENCE, ROOT, SQL_FILES, print_exception
from trioexplorer.lib.xtable import XTable
from trioexplorer.lib.pagedresults import PagedResults
import pydot
import sys
import os
from pylons import config
from paste.urlparser import StaticURLParser

log = logging.getLogger(__name__)

__id__ = "$Id$"

def needsLogin(fn):
    '''Login decorator: if the user is not logged in, send up a login
    rather than the page content.  Point the login form back to the
    same page to avoid a redirect when the user logs in correctly.
    '''
    def _wrapper(*args, **kwargs):
        if 'user' in session and cnx.has_key(session['user'].name):
            # User is logged in; allow access
            #return fn(*args, **kwargs)
            return fn(*args)
        else:
            return redirect_to('login')
    return _wrapper

def addToHistory(query):
    if query.lstrip()[0:4].lower() != "view" and query.find('"') == -1:
        query = query.replace('\r', ' ').replace('\n', ' ')
        query = query.replace('\\', '\\\\').replace("'", "\\'")
        if session['user'].history.count(query) != 0:
            session['user'].history.remove(query)
        session['user'].history.append(query)
        if len(session['user'].history) == session['user'].history_size + 1:
            del session['user'].history[0]
        session.save()

def calc_conf(table, aid, disj='0'):
    tcnx = cnx[session['user'].name]
    tnc = tcnx.ncursor()
    tables = table.split(",")
    aids = [int(x) for x in aid.split(",")]
    disjs = [int(x) for x in disj.split(",")]

    log.debug(session['user'].name + ": " + str(tables) + "|" + str(aids))

    query = "SELECT conf(array[%s], array[%s], array[%s], NULL, 0, 1)" \
            % (", ".join("'%s'" % x for x in tables),
               ", ".join("%d" % x for x in aids),
               ", ".join("%d" % x for x in disjs))
    conf = None

    # Memoize result
    try:
        tnc.execute(query)
        conf = tnc.fetchone()[0]
        tnc.close()
        tcnx.commit()
    except Exception, ex:
        tcnx.rollback()
        conf = "error"
        log.warning(ex)

    log.debug(session['user'].name + ": done")
    return conf

def executeQuery(query, limit=-1, num_results=-1):
    res_lxtps = []
    error = None
    try:
        tcnx = cnx[session['user'].name]
        tc = tcnx.cursor()
        for q in query.split(";"):
            if q != '' and not q.isspace():
                tc.execute(q)

        conf_computed = True
        if tc.isFetchable:
            while True:
                try:
                    xtuple = tc.xfetchone()
                    if xtuple is None or (limit != -1 and len(res_lxtps) == limit):
                        break
                    conf_computed = conf_computed and xtuple.getConfidence(tc.version) is not None
                    res_lxtps.append(xtuple)
                except Exception, ex:
                    break

            if not conf_computed and not tc.gotAsConf:
                conf_computed = False
            elif not conf_computed and tc.tableType == TYPE_CONFIDENCE and len(res_lxtps) > 0:
                lin = res_lxtps[0][0].getLineage()
                if len(lin.srctable) == 1:
                    tnc = cnx[session['user'].name].ncursor()
                    q = "SELECT confbase FROM trio_tables WHERE tablename = '%s'" % lin.srctable[0]
                    try:
                        tnc.execute(q)
                        conf_computed = tnc.fetchone()[0]
                        tnc.close()
                    except Exception, ex:
                        conf_computed = False
        tcnx.commit()

    except Exception, error:
        try:
            tcnx.rollback()
        except Exception, ex:
            log.warning(ex)
        print_exception()
        return PagedResults(query, res_lxtps, error, tc.description, tc.tableType)

    return PagedResults(query, res_lxtps, error, tc.description, tc.tableType, tc.isFetchable, tc.nolineage, conf_computed, tc.isViewTable, tc.distinct, tc.u_columns, tc.version, num_results, tc.extLineage)

def generateGraph():
    """Generate relation graph."""
    tcnx = cnx[session['user'].name]
    tc = tcnx.cursor()
    tc.execute("\d")

    graph = pydot.Dot(graph_name="schema")
    for xr in tc.fetchall():
        thisRelation = xr[0].upper()
        tname =  xr[0]
        xt = XTable(tname, cnx[session['user'].name])
        for attr in xt.getAttributes():
            thisRelation += "|" + str(attr[0]).strip("[()]*\'\ ").lower()
            if not attr[1]:
                thisRelation += "*"
        node = pydot.Node(tname, label="{" + thisRelation + "}")

        type = xt.getType()
        if type == TYPE_CONFIDENCE:
            node.set_color("orange")
        elif type == TYPE_UNCERTAIN:
            node.set_color("green")
        else:
            node.set_color("navy")

        node.set_fontname("Sans")
        #node.set_fontname("Times-Roman")
        node.set_fontsize("12")
        node.set_shape("Mrecord")
        node.set_URL("index?query=view table %s;" % tname)

        l = xt.getLineage()
        for parent in l:
            edge = pydot.Edge(tname, parent)
            edge.set_color("skyblue2")
            graph.add_edge(edge)
        if l == []:
            node.set_shape("record")
        graph.add_node(node)

    tcnx.commit()

    graph.write_png(path="%s/graphs/graph.%s.png"
                         % (config['pylons.cache_dir'], session['user'].name))
    graph.write_cmapx(path="%s/graphs/graph.%s.cmapx"
                           % (config['pylons.cache_dir'], session['user'].name))

class TrioUserData:
    def __init__(self, username):
        self.name = username

        # Hold query history
        self.history = []

        # Number of queries saved in self.history.
        self.history_size = 10

        # Hold queries loaded from file.
        self.queries = []

        # Hold result of a query.
        self.result = None

        # Number of xtuples displayed in one page.
        self.num_tuples = 10

        # Used to assign unique id to alternatives.
        self.sequence = [0]

cnx = {} # Holds trio connection(s).
class TrioxplorerController(BaseController):
    def __init__(self):
        pass

    #
    # Pages before login
    #
    def createDB(self):
        return render('/createdb.mako')

    def doCreateDB(self):
        user = request.params['user']
        password = request.params['password']
        password2 = request.params['password2']

        if password != password2:
            session['error']  = 'Your passwords did not match'
            session.save()
            return h.redirect_to('createDB')

        if user == '' or password ==''or password2 =='':
            session['error'] = 'Required field cannot be left blank'
            session.save()
            return h.redirect_to('createDB')

        try:
            db = None
            conf = config['app_conf']

            # Create a new database (as DBA)
            if "username" in conf and "password" in conf:
                db = pgdb.connect(database=conf["username"], user=conf["username"], password=conf["password"])
            elif "default_user" in conf:
                db = pgdb.connect(database=conf["default_user"], user=conf["default_user"])
            else:
                db = pgdb.connect()
            c = db.cursor()
            c.execute('COMMIT')
            c.execute('CREATE DATABASE "%s"' % user)
            c.execute('CREATE ROLE "%s" LOGIN PASSWORD \'%s\'' % (user, password))
            c.execute('ALTER DATABASE "%s" OWNER TO "%s"' %  (user, user))
            c.close()
            db.close()

            # Set up stored procedures (as DBA)
            if "username" in conf and "password" in conf:
                db = pgdb.connect(database=user, user=conf["username"], password=conf["password"])
            elif "default_user" in conf:
                db = pgdb.connect(database=user, user=conf["default_user"])
            else:
                db = pgdb.connect(database=user)
            c = db.cursor()
            for sql_file in SQL_FILES:
                file = open(sql_file, "r")
                q = file.read()
                c.execute(q)
            c.close()
            db.commit()
            db.close()

            # Set up Trio system catalogs (as user)
            db = pgdb.connect(user=user, database=user, password=password)
            c = db.cursor()
            file = open(os.path.join(ROOT, 'sql', 'setup.sql'), "r")
            q = file.read()
            c.execute(q)
            c.close()
            db.commit()
            db.close()

        except Exception, ex:
            if 'database "' in unicode(ex) and '" already exists' in unicode(ex):
                session['error'] = 'User already exists'
            else:
                session['error'] = unicode(ex)
            session.save()
            return h.redirect_to('createDB')

        return self.doLogin(user, password)

    def deleteDB(self):
        try:
            cnx[session['user'].name].close()
            del cnx[session['user'].name]
            db = None
            conf = config['app_conf']
            if "username" in conf and "password" in conf:
                db = pgdb.connect(user=conf["username"], password=conf["password"])
            elif "default_user" in conf:
                db = pgdb.connect(database=conf["default_user"], user=conf["default_user"])
            else:
                db = pgdb.connect()
            c = db.cursor()
            c.execute('COMMIT')
            c.execute('DROP DATABASE "%s"' % session['user'].name)
            c.execute('DROP ROLE "%s"' % session['user'].name)
            c.close()
            db.close()
        except Exception, ex:
            session['error'] = unicode(ex)
            session.save()

        return h.redirect_to('login')

    def doLogin(self, user='', password=''):
        if user == "" and password == "":
            user = request.params['user']
            password = request.params['password']

        if user == "" or password == "":
            session['error'] = 'Required field cannot be left blank'
            session.save()
            return h.redirect_to('login')

        try:
            tdb = triodb.connect(user=user, database=user, password=password)
        except Exception, ex:
            if 'password authentication failed' in unicode(ex):
                session['error'] = "Username and password didn't match"
            elif 'does not exist' in unicode(ex):
                session['error'] = "Username doesn't exist"
            else:
                session['error'] = unicode(ex)
            session.save()
            return h.redirect_to('login')

        session['user'] = TrioUserData(user)
        cnx[session['user'].name] = tdb
        session.save()
        return h.redirect_to('index')

    def login(self):
        return render("/login.mako")

    def logout(self):
        try:
            cnx[session['user'].name].close()
            del cnx[session['user'].name]
            del session['user']
        except:
            pass
        session.save()
        return h.redirect_to('login')

    def clearHistory(self):
        session['user'].history= []
        session.save()
        return h.redirect_to('index')

    #
    # Pages after login
    #
    @needsLogin
    def index(self, compconf=False):
        c.tcnx = cnx[session['user'].name]
        query = ''
        page_num = -1
        execute  = 'True'
        compconf = 'False'
        xindex = -1

        if request.params.has_key('query'):
            query = request.params['query']
        if request.params.has_key('page_num'):
            page_num = request.params['page_num']
        if request.params.has_key('execute'):
            execute = request.params['execute']
        if request.params.has_key('compconf'):
            compconf = request.params['compconf']
        if request.params.has_key('xindex'):
            xindex = request.params['xindex']

        c.query = query
        c.results = None
        num_results = 0

        log.info(session['user'].name + ": " + query)
        log.info(session['user'].name + ": compconf=" + compconf + ", execute=" + execute)

        if query.lstrip()[0:4].lower() == 'view':
            try:
                tcnx = cnx[session['user'].name]
                tnc = tcnx.ncursor()
                tnc.execute("SELECT count(distinct trio_xid) FROM %s" % str(query).split(" ")[-1])
                num_results = tnc.fetchone()[0]
                tnc.close()
                tcnx.commit()
            except Exception, error:
                e = sys.exc_info()
                log.debug(error)
                try:
                    tcnx.rollback()
                except Exception, ex:
                    log.warning(ex)

        if num_results > 100:
            if page_num == -1:
                addToHistory(query)
                c.results = executeQuery(query, 10, num_results)
                session['user'].results = c.results
                session.save()
            elif page_num >= 0:
                c.results = executeQuery(query, (int(page_num) + 1) * 10, num_results)
                session['user'].results = c.results
                session.save()
                c.results.setPage(int(page_num))
        else:
            if query != '' and not query.isspace() and execute != 'False' and page_num == -1:
                addToHistory(query)
                c.results = executeQuery(query)
                session['user'].results = c.results
                if c.results.error is None:
                    session.save()
            elif page_num > 0:
                c.results = session['user'].results
                session['user'].results.setPage(int(page_num))
                session.save()

        if compconf == 'True':
            if xindex >= 0:
                xtup = session['user'].results.results[int(xindex)]
                aindex = 0
                for alt in xtup:
                    table = ",".join(alt.lineage.tables())
                    aid = ",".join([str(x) for x in alt.lineage.aids()])
                    disj = ",".join([str(x) for x in alt.lineage.disj()])
                    conf = calc_conf(table, aid, disj)
                    session['user'].results.results[int(xindex)][int(aindex)].confidence = conf
                    aindex += 1
            else:
                xindex = 0
                session['user'].results.conf_computed = True
                for xtup in session['user'].results.results:
                    aindex = 0
                    for alt in xtup:
                        table = ",".join(alt.lineage.tables())
                        aid = ",".join([str(x) for x in alt.lineage.aids()])
                        disj = ",".join([str(x) for x in alt.lineage.disj()])
                        conf = calc_conf(table, aid, disj)
                        session['user'].results.results[int(xindex)][int(aindex)].confidence = conf
                        aindex += 1
                    xindex += 1
            session.save()

        log.info(session['user'].name + ": done\n")
        return render("/index.mako")

    @needsLogin
    def abort(self):
        cnx[session['user'].name].pdb._pgdbCnx__cnx.cancel()
        return  Response("")

    def confidence(self):
        table = request.params['table']
        aid = request.params['aid']
        try:
            disj = request.params['disj']
        except:
            disj = "0"
        xindex, aindex = -1, -1
        if request.params.has_key('xindex'):
            xindex = request.params['xindex']
        if request.params.has_key('aindex'):
            aindex = request.params['aindex']

        conf = calc_conf(table, aid, disj)
        if xindex >= 0 and aindex >= 0:
            session['user'].results.results[int(xindex)][int(aindex)].confidence = conf
            session.save()

        if isinstance(conf, float):
            return Response("%.3f" % conf)
        else:
            return Response(str(conf))

    @needsLogin
    def lineage(self):
        c.tcnx = cnx[session['user'].name]
        c.top = False
        if request.params.has_key('top'):
            c.top = True
        if request.params.has_key('tables'):
            c.tables = request.params['tables'].split(',')
        if request.params.has_key('aids'):
            c.aids = request.params['aids'].split(',')
        if request.params.has_key('disj'):
            c.disj = request.params['disj'].split(',')
            c.disj = [int(x) for x in c.disj]
        return render("/lineage.mako")

    @needsLogin
    def extlineage(self):
        c.tcnx = cnx[session['user'].name]
        if request.params.has_key('tables'):
            c.tables = request.params['tables'].split(',')
        if request.params.has_key('aids'):
            c.aids = request.params['aids'].split(',')
        return render("/extlineage.mako")

    @needsLogin
    def graph(self):
        static = StaticURLParser('%s/graphs' % config['pylons.cache_dir'])
        request.environ['PATH_INFO'] = '/graph.%s.png' % session['user'].name
        return static(request.environ, self.start_response)

    @needsLogin
    def loadData(self):
        tcnx = cnx[session['user'].name]
        tc = tcnx.cursor()
        tc.execute("\d")
        tables = tc.fetchall()
        tcnx.commit()

        for table in tables:
            executeQuery("Drop Table %s Cascade" % table[0])

        if request.params.has_key('dataset'):
            dataset = request.params['dataset']
        else:
            dataset = 'crime'

        if dataset == 'movie':
            fname = "%s/triql/movie_data.triql" % config['pylons.paths']['static_files']
        elif dataset == 'product':
            fname = "%s/triql/product_data.triql" % config['pylons.paths']['static_files']
        else:
            fname = "%s/triql/crime_data.triql" % config['pylons.paths']['static_files']

        f = open(fname)
        try:
            for line in f:
                tc.execute(line.rstrip('; \t\n\r'))
            tcnx.commit()
            session['loaded'] = True
        except:
            tcnx.rollback()
            session['loaded'] = False

        session.save()
        return h.redirect_to('samples')

    @needsLogin
    def upload(self):
        session['user'].queries = []
        query = ""
        filedata = request.params['myFile']
        if filedata == '':
            return h.redirect_to('scripts')
        for line in filedata.file.readlines():
            if line[:-1].strip == "":
                continue
            if line.startswith("--"):
                session['user'].queries.append(line)
            else:
                if query == "":
                    query = line
                else:
                    query += " " + line
                if query[:-1].strip().endswith(";"):
                    session['user'].queries.append(query.strip())
                    query = ""
        session.save()
        return h.redirect_to('scripts')

    @needsLogin
    def schema(self):
        generateGraph()
        return render("/schema.mako")

    @needsLogin
    def samples(self):
        c.tcnx = cnx[session['user'].name]
        return render("/samples.mako")

    @needsLogin
    def scripts(self):
        return render("/scripts.mako")

    @needsLogin
    def help(self):
        return render("/help.mako")
