#!/usr/bin/python


# Standard Library
import inspect
import json
import logging
import os
import sys
import time
import traceback
import uuid
from collections import defaultdict
from Queue import Queue
from cStringIO import StringIO
from threading import Thread, Lock

# Third Party
import boto.s3
import mixingboard
import pip

# Local
from chassis.database import db_session
from chassis.models import JobHistory
from chassis.util import processFormattedTableDescription
from receiver import EventBroadcastReceiver


# import pyspark junk (REQUIRES SPARK TO BE ON SYSPATH)
from pyspark import SparkFiles
from pyshark import SharkConf
from sparkler import SparklerContext


# set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# defaults
WORKER_MEMORY = "3300m"
SCHEDULER_POOLS = [
    "verylow",
    "low",
    "default",
    "high",
    "veryhigh"
]


class JobRunner:
    """
    Manages and runs pyspark jobs
    """

    MAX_FILE_SIZE = 1024*1024*16 # don't write out results files larger than 16 MBs

    def __init__(self, account, master, uploadFolder, iamUsername,
                accessKeyId, accessKeySecret, conf={}):
        """
        Initializes a job runner

        Params:
            account: a user account for this job runner
            conf: allow you to set configuration variables on a spark context
        """

        self.account = account
        self.region = mixingboard.REGION
        self.master = master
        self.conf = conf
        self.uploadFolder = uploadFolder
        self.iamUsername = iamUsername
        self.accessKeyId = accessKeyId
        self.accessKeySecret = accessKeySecret

        self.handles = defaultdict(dict)
        self.handlesLock = Lock()

        self.createContext()

    # TODO when the runner is isolated, this needs to pass back the history over stdout
    def _makeHistory(self, accountId, userId, event, jobId=None, jobHandle=None, jobType="spark", data={}):
        """
        Save an entry to the job history record

        Params:
            account: an account
            user: a user
            event: an event type to save
            jobId: a job id
            data: a json serializable object
        Returns:
            a history event
        """

        jobHistory = JobHistory(account_id=accountId, user_id=userId, event=event, jobType=jobType,
                                jobId=jobId, jobHandle=jobHandle, data=data)
        db_session.add(jobHistory)
        db_session.commit()

        return jobHistory

    def _makeHandle(self):
        """
        Generates a unique query handle. Not completely resistant to
        collisions but it's nearly impossible so... close enough

        Returns:
            a unique handle
        """

        return "%s_%s" % (
            uuid.uuid4().bytes.encode("base64")[:21].translate(None, "/+"),
            int(time.time()*1000)
        )

    def runSQLSync(self, sql):
        """
        Run a sql query and wait for the result

        Params:
            sql: a sql query
        Returns:
            the result of running a job
        """

        result = [row.split('\t') for row in self.context.sql(sql)]

        return result


    def _parseSQL(self, sql):

        # if we have multiple queries, break them apart
        queries = []
        buff = ""
        skip = False
        parenDepth = 0
        quoteType = None

        for c in sql:

            buff += c

            # deal with escaping
            if c == "\\":
                skip = True
                continue

            if not skip:

                if quoteType is None:

                    if c == "(":
                        parenDepth += 1

                    elif c == ")" and parenDepth > 0:
                        parenDepth -= 1

                    elif c == "'" or c == '"':
                        quoteType = c

                    elif c == ";" and parenDepth == 0:
                        buff = buff.strip()
                        queries.append(buff[:-1])
                        buff = ""
                        continue

                elif quoteType == c:
                    quoteType = None

            skip = False

        buff = buff.strip()
        if len(buff) > 0:
            queries.append(buff)

        return queries


    def runSQLAsync(self, sql, options={}):
        """
        Run a job

        Params:
            sql: a sql query
        Returns:
            a handle for retrieveing future information about a job
        """

        handle = self._makeHandle()

        def sqlRunInner(sql, handle, options):

            self.handlesLock.acquire()
            self.handles[handle]['running'] = True
            self.handles[handle]['options'] = options
            self.handlesLock.release()

            result = None
            queryResult = None
            columns = None
            outputDir = '/tmp/%s/shark/%s' % (
                self.iamUsername,
                handle,
            )

            self._makeHistory(
                self.account,
                options.get('user'),
                'job_start',
                jobHandle=handle,
                data={ "jobName": options.get("jobName") },
                jobType="sql"
            )


            try:

                self.context.setJobGroup(handle, "SQL Query by Jerb")
                self.context.setLocalProperty("spark.job.name", options.get("jobName","Untitled"))
                self.context.setLocalProperty("spark.job.type", options.get("jobType","sql"))

                poolName = options.get("priority","default")
                pool = poolName if poolName in SCHEDULER_POOLS else "default"
                self.context.setLocalProperty("spark.scheduler.pool", options.get("priority","default"))

                queries = self._parseSQL(sql)

                self.context.setLocalProperty('spark.job.numqueries', str(len(queries)))

                for i, sql in enumerate(queries):

                    # if we're at the last query, save it's output
                    if i == len(queries)-1:

                        try:
                            # TODO use hiveContext stuff and saveAsTextFile
                            self.context.sql("""
                                CREATE TABLE `__%s`