# Licensed to the StackStorm, Inc ('StackStorm') under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import itertools

import six

from st2api.controllers.resource import ResourceController
from st2common import log as logging
from st2common.constants import action as action_constants
from st2common.models.db.execution import ActionExecutionOutputDB
from st2common.models.api.execution import ActionExecutionAPI
from st2common.models.api.execution import ActionExecutionOutputAPI
from st2common.persistence.execution import ActionExecution
from st2common.persistence.execution import ActionExecutionOutput
from st2common.router import Response
from st2common.util.jsonify import json_encode
from st2common.rbac.types import PermissionType
from st2common.stream.listener import get_listener

__all__ = [
    'ActionExecutionOutputStreamController'
]

LOG = logging.getLogger(__name__)


class ActionExecutionOutputStreamController(ResourceController):
    model = ActionExecutionAPI
    access = ActionExecution

    supported_filters = {
        'output_type': 'output_type'
    }

    CLOSE_STREAM_LIVEACTION_STATES = action_constants.LIVEACTION_COMPLETED_STATES + [
        action_constants.LIVEACTION_STATUS_PAUSING,
        action_constants.LIVEACTION_STATUS_RESUMING
    ]

    def get_one(self, id, output_type='all', requester_user=None):
        # Special case for id == "last"
        if id == 'last':
            execution_db = ActionExecution.query().order_by('-id').limit(1).first()

            if not execution_db:
                raise ValueError('No executions found in the database')

            id = str(execution_db.id)

        execution_db = self._get_one_by_id(id=id, requester_user=requester_user,
                                           permission_type=PermissionType.EXECUTION_VIEW)
        execution_id = str(execution_db.id)

        query_filters = {}
        if output_type and output_type != 'all':
            query_filters['output_type'] = output_type

        def format_output_object(output_db_or_api):
            if isinstance(output_db_or_api, ActionExecutionOutputDB):
                data = ActionExecutionOutputAPI.from_model(output_db_or_api)
            elif isinstance(output_db_or_api, ActionExecutionOutputAPI):
                data = output_db_or_api
            else:
                raise ValueError('Unsupported format: %s' % (type(output_db_or_api)))

            event = 'st2.execution.output__create'
            result = 'event: %s\ndata: %s\n\n' % (event, json_encode(data, indent=None))
            return result

        def existing_output_iter():
            # Consume and return all of the existing lines
            output_dbs = ActionExecutionOutput.query(execution_id=execution_id, **query_filters)

            # Note: We return all at once instead of yield line by line to avoid multiple socket
            # writes and to achieve better performance
            output = [format_output_object(output_db) for output_db in output_dbs]
            output = ''.join(output)
            yield six.binary_type(output.encode('utf-8'))

        def new_output_iter():
            def noop_gen():
                yield six.binary_type(''.encode('utf-8'))

            # Bail out if execution has already completed / been paused
            if execution_db.status in self.CLOSE_STREAM_LIVEACTION_STATES:
                return noop_gen()

            # Wait for and return any new line which may come in
            execution_ids = [execution_id]
            listener = get_listener(name='execution_output')  # pylint: disable=no-member
            gen = listener.generator(execution_ids=execution_ids)

            def format(gen):
                for pack in gen:
                    if not pack:
                        continue
                    else:
                        (_, model_api) = pack

                        # Note: gunicorn wsgi handler expect bytes, not unicode
                        # pylint: disable=no-member
                        if isinstance(model_api, ActionExecutionOutputAPI):
                            if output_type and output_type != 'all' and \
                               model_api.output_type != output_type:
                                continue

                            output = format_output_object(model_api).encode('utf-8')
                            yield six.binary_type(output)
                        elif isinstance(model_api, ActionExecutionAPI):
                            if model_api.status in self.CLOSE_STREAM_LIVEACTION_STATES:
                                yield six.binary_type(''.encode('utf-8'))
                                break
                        else:
                            LOG.debug('Unrecognized message type: %s' % (model_api))

            gen = format(gen)
            return gen

        def make_response():
            app_iter = itertools.chain(existing_output_iter(), new_output_iter())
            res = Response(content_type='text/event-stream', app_iter=app_iter)
            return res

        res = make_response()
        return res


action_execution_output_controller = ActionExecutionOutputStreamController()

@@ -1559,16 +1559,185 @@
 ame__)%0A%0A
+# Event which is returned when no more data will be produced on this stream endpoint before closing%0A# the connection.%0ANO_MORE_DATA_EVENT = 'event: EOF%5Cndata: %5C'%5C'%5Cn%5Cn'%0A%0A
 %0Aclass A
@@ -4029,34 +4029,50 @@
 six.binary_type(
-''
+NO_MORE_DATA_EVENT
 .encode('utf-8')
@@ -5472,18 +5472,34 @@
 ry_type(
-''
+NO_MORE_DATA_EVENT
 .encode(

Send a special EOF even when there will be no more events on a particual stream endpoint before closing the connection.