import logging
import select
import subprocess

import luigi
import time

from luigi_gcloud.gcore import get_default_client, _GCloudTask
from luigi_gcloud.storage import GCSFileSystem

logger = logging.getLogger('luigi-gcloud')


class _DataProcJob:
    def __init__(self, dataproc_api, project_id, job):
        self.dataproc_api = dataproc_api
        self.project_id = project_id
        self.job = dataproc_api.projects().jobs().submit(
                projectId=self.project_id,
                body=job).execute()
        self.job_id = self.job['reference']['jobId']
        logger.info('DataProc job %s is %s', self.job_id, str(self.job['status']['state']))

    def wait_for_done(self):
        while True:
            self.job = self.dataproc_api.projects().jobs().get(
                    projectId=self.project_id,
                    jobId=self.job_id).execute()
            if 'ERROR' == self.job['status']['state']:
                print(str(self.job))
                logger.error('DataProc job %s has errors', self.job_id)
                logger.error(self.job['status']['details'])
                logger.debug(str(self.job))
                return False
            if 'DONE' == self.job['status']['state']:
                return True
            logger.info('DataProc job %s is %s', self.job_id, str(self.job['status']['state']))
            time.sleep(5)

    def raise_error(self, message=None):
        if 'ERROR' == self.job['status']['state']:
            if message is None:
                message = "Google DataProc job has error"
            raise Exception(message + ": " + str(self.job['status']['details']))

    def get(self):
        return self.job


class DataProcPigTask(_GCloudTask):
    service_name = "dataproc"

    def query_file(self):
        return None

    def query_uri(self):
        return self.client.project_staging() + self.resolved_name() + ".pig"

    def lib_uris(self):
        return []

    def run(self):
        http = self.client.http_authorized()
        dataproc_api = self.client.dataproc_api(http)
        name = self.resolved_name()

        if self.query_file() is not None:
            fs = GCSFileSystem()
            fs.put(self.query_file(),
                   self.query_uri())

        artifacts = self.lib_uris()

        job = {
            "job": {
                "reference": {
                    "projectId": self.client.project_id(),
                    "jobId": name,
                },
                "placement": {
                    "clusterName": self.get_service_value("clusterName", "cluster-1")
                },
                "pigJob": {
                    "continueOnFailure": False,
                    "queryFileUri": self.query_uri(),
                    "jarFileUris": artifacts,
                }
            }
        }
        variables = self.variables()
        if variables is not None:
            job["job"]["pigJob"]["scriptVariables"] = variables

        print(job)
        submitted = _DataProcJob(dataproc_api, self.client.project_id(), job)
        if not submitted.wait_for_done():
            submitted.raise_error("DataProcTask has errors")


class DataProcSparkTask(_GCloudTask):
    service_name = "dataproc"

    def job_file(self):
        return None

    def job_uri(self):
        return self.client.project_staging() + self.resolved_name() + ".jar"

    def lib_uris(self):
        return []

    def main(self):
        raise NotImplementedError

    def args(self):
        return []

    def run(self):
        http = self.client.http_authorized()
        dataproc_api = self.client.dataproc_api(http)
        name = self.resolved_name()

        if self.job_file() is not None:
            logger.warning("Copying job artifact to storage. Consider uploading outside of job.")
            fs = GCSFileSystem()
            fs.put(self.get_service_value("basePath", ".") + self.job_file(),
                   self.job_uri())

        artifacts = [self.job_uri()]
        artifacts.extend(self.lib_uris())

        job = {
            "job": {
                "reference": {
                    "projectId": self.client.project_id(),
                    "jobId": name,
                },
                "placement": {
                    "clusterName": self.get_service_value("clusterName", "cluster-1")
                },
                "sparkJob": {
                    "mainClass": self.main(),
                    "jarFileUris": artifacts,
                }
            }
        }
        args = self.args()
        if args is not None:
            job["job"]["sparkJob"]["args"] = args

        print(job)
        submitted = _DataProcJob(dataproc_api, self.client.project_id(), job)
        if not submitted.wait_for_done():
            submitted.raise_error("DataProcTask has errors")

@@ -12,54 +12,8 @@
 ing%0A
-import select%0Aimport subprocess%0A%0Aimport luigi%0A
 impo
@@ -55,28 +55,8 @@
 port
- get_default_client,
  _GC
@@ -156,16 +156,40 @@
 loud')%0A%0A
+default_pig_udf = None%0A%0A
 %0Aclass _
@@ -3457,32 +3457,471 @@
      return %5B%5D%0A%0A
+    def _jar_file_uris(self):%0A        global default_pig_udf%0A        artifacts = %5Bself.job_uri()%5D%0A        uris = self.lib_uris()%0A        artifacts.extend(uris)%0A        if default_pig_udf is not None:%0A            if not default_pig_udf.append and len(uris) == 0:%0A                artifacts.extend(default_pig_udf.uris)%0A            elif default_pig_udf.append:%0A                artifacts.extend(default_pig_udf.uris)%0A        return artifacts%0A%0A
     def run(self
@@ -3915,32 +3915,32 @@
  def run(self):%0A
-
         http = s
@@ -4348,88 +4348,8 @@
 ))%0A%0A
-        artifacts = %5Bself.job_uri()%5D%0A        artifacts.extend(self.lib_uris())%0A%0A
     
@@ -4768,33 +4768,45 @@
 rFileUris%22: 
-artifacts
+self._jar_file_uris()
 ,%0A          
@@ -5075,32 +5075,32 @@
 ait_for_done():%0A
-
             subm
@@ -5124,28 +5124,179 @@
 (%22DataProcTask has errors%22)%0A
+%0A%0Adef set_default_pig_udf(uris, append=False):%0A    global default_pig_udf%0A    default_pig_udf = %7B%0A        'append': append,%0A        'uris': uris%0A    %7D%0A

Allow setting default UDF's for pig closes #18