from sqlalchemy.orm import Session
from aiohttp import web
import numpy as np
import asyncio
import logging
from joule import app_keys

from joule.models import (folder, DataStore, DataStream,
                          InsufficientDecimationError, DataError,
                          pipes)
from joule.models.supervisor import Supervisor

from joule.errors import SubscriptionError

log = logging.getLogger('joule')


async def read_json(request: web.Request):
    return await read(request, json=True)


async def read(request: web.Request, json=False):
    if 'subscribe' in request.query and request.query['subscribe'] == '1':
        return await _subscribe(request, json)
    else:
        return await _read(request, json)


async def _read(request: web.Request, json):
    db: Session = request.app[app_keys.db]
    data_store: DataStore = request.app[app_keys.data_store]
    # find the requested stream
    if 'path' in request.query:
        stream = folder.find_stream_by_path(request.query['path'], db, stream_type=DataStream)
    elif 'id' in request.query:
        stream = db.get(DataStream,request.query["id"])
    else:
        return web.Response(text="specify an id or a path", status=400)
    if stream is None:
        return web.Response(text="stream does not exist", status=404)

    # parse optional parameters
    params = {'start': None, 'end': None, 'max-rows': None, 'decimation-level': None}
    param = ""  # to appease type checker

    try:
        for param in params:
            if param in request.query:
                params[param] = int(request.query[param])
    except ValueError:
        return web.Response(text="parameter [%s] must be an int" % param, status=400)

    # make sure parameters make sense
    if ((params['start'] is not None and params['end'] is not None) and
            (params['start'] >= params['end'])):
        return web.Response(text="[start] must be < [end]", status=400)
    if params['max-rows'] is not None and params['max-rows'] <= 0:
        return web.Response(text="[max-rows] must be > 0", status=400)
    if params['decimation-level'] is not None and params['decimation-level'] <= 0:
        return web.Response(text="[decimation-level] must be > 0", status=400)

    # --- Binary Streaming Handler ---
    resp = None

    async def stream_data(data: np.ndarray, layout, factor):
        nonlocal resp

        if resp is None:
            resp = web.StreamResponse(status=200,
                                      headers={'joule-layout': layout,
                                               'joule-decimation': str(factor)})
            resp.enable_chunked_encoding()
            await resp.prepare(request)
        await resp.write(data.tobytes())

    # --- JSON Handler ---

    data_blocks = []  # array of data segments
    data_segment = None
    decimation_factor = 1

    async def retrieve_data(data: np.ndarray, layout, factor):
        nonlocal data_blocks, data_segment, decimation_factor
        decimation_factor = factor
        if np.array_equal(data, pipes.interval_token(layout)):
            if data_segment is not None:
                data_blocks.append(data_segment.tolist())
                data_segment = None
        else:
            data = np.c_[data['timestamp'][:, None], data['data']]

            if data_segment is None:
                data_segment = data
            else:
                data_segment = np.vstack((data_segment, data))

    if json:
        callback = retrieve_data
    else:
        callback = stream_data

    # create an extraction task
    try:
        await data_store.extract(stream, params['start'], params['end'],
                                 callback=callback,
                                 max_rows=params['max-rows'],
                                 decimation_level=params['decimation-level'])
    except InsufficientDecimationError as e:
        return web.Response(text="decimated data is not available: %s" % e, status=400)
    except DataError as e:
        msg = str(e)
        if 'no such stream' in msg.lower() and (params['decimation-level'] is not None):  # pragma: no cover
            # clean up error message when user requested a particular decimation level
            msg = "requested decimation level [%d] does not exist" % params['decimation-level']
        return web.Response(text="read error: %s" % msg, status=400)

    if json:
        # put the last data_segment on
        if data_segment is not None:
            data_blocks.append(data_segment.tolist())
        return web.json_response({"data": data_blocks, "decimation_factor": decimation_factor})
    else:
        if resp is None:
            return web.json_response(text="this stream has no data", status=400)
        return resp


async def _subscribe(request: web.Request, json: bool):
    db: Session = request.app[app_keys.db]
    supervisor: Supervisor = request.app[app_keys.supervisor]
    if json:
        return web.Response(text="JSON subscription not implemented", status=400)

    # find the requested stream
    if 'path' in request.query:
        stream = folder.find_stream_by_path(request.query['path'], db, stream_type=DataStream)
    elif 'id' in request.query:
        stream = db.get(DataStream,request.query["id"])
    else:
        return web.Response(text="specify an id or a path", status=400)
    if stream is None:
        return web.Response(text="stream does not exist", status=404)
    pipe = pipes.LocalPipe(stream.layout)
    try:
        unsubscribe = supervisor.subscribe(stream, pipe)
    except SubscriptionError:
        return web.Response(text="stream is not being produced", status=400)
    resp = web.StreamResponse(status=200,
                              headers={'joule-layout': stream.layout,
                                       'joule-decimation': '1'})
    resp.enable_chunked_encoding()

    try:
        await resp.prepare(request)
    except ConnectionResetError:
        unsubscribe()
        return resp

    try:
        while True:
            try:
                data = await pipe.read()
            except pipes.EmptyPipe:
                unsubscribe()
                return resp
            pipe.consume(len(data))
            if len(data) > 0:
                await resp.write(data.tobytes())
            if pipe.end_of_interval:
                await resp.write(pipes.interval_token(stream.layout).tobytes())
    except asyncio.CancelledError as e:
        unsubscribe()
        # propogate the CancelledError up
        raise e
    except ConnectionResetError:
        unsubscribe()
        return resp


async def intervals(request: web.Request):
    db: Session = request.app[app_keys.db]
    data_store: DataStore = request.app[app_keys.data_store]
    # find the requested stream
    if 'path' in request.query:
        stream = folder.find_stream_by_path(request.query['path'], db, stream_type=DataStream)
    elif 'id' in request.query:
        stream = db.get(DataStream,request.query["id"])
    else:
        return web.Response(text="specify an id or a path", status=400)
    if stream is None:
        return web.Response(text="stream does not exist", status=404)
    # parse time bounds if specified
    try:
        if 'start' in request.query:
            start = int(request.query['start'])
        else:
            start = None
        if 'end' in request.query:
            end = int(request.query['end'])
        else:
            end = None
    except ValueError:
        return web.Response(text="[start] and [end] must be an integers", status=400)

    # make sure parameters make sense
    if (start is not None and end is not None) and start >= end:
        return web.Response(text="[start] must be < [end]", status=400)

    return web.json_response(await data_store.intervals(stream, start, end))


async def write(request: web.Request):
    db: Session = request.app[app_keys.db]
    data_store: DataStore = request.app[app_keys.data_store]
    # find the requested stream
    if 'path' in request.query:
        stream = folder.find_stream_by_path(request.query['path'], db, stream_type=DataStream)
    elif 'id' in request.query:
        stream = db.get(DataStream,request.query["id"])
    else:
        return web.Response(text="specify an id or a path", status=400)
    if stream is None:
        return web.Response(text="stream does not exist", status=404)
    # spawn in inserter task
    stream.is_destination = True
    db.commit()

    pipe = pipes.InputPipe(name="inbound", stream=stream, reader=request.content)
    if 'merge-gap' in request.query:
        try:
            merge_gap = int(request.query['merge-gap'])
        except ValueError:
            return web.Response(text="merge-gap must be an integer", status=400)
        
        if merge_gap < 0:
            return web.Response(text="merge-gap must be >= 0", status=400)
    else:
        merge_gap = 0
    try:
        task = await data_store.spawn_inserter(stream, pipe, insert_period=0, merge_gap=merge_gap)
        await task
    except DataError as e:
        stream.is_destination = False
        db.commit()
        print("closing stream due to error")
        return web.Response(text=str(e), status=400)
    except asyncio.CancelledError as e:
        raise e
    finally:
        stream.is_destination = False
        db.commit()
    return web.Response(text="ok")


async def remove(request: web.Request):
    db: Session = request.app[app_keys.db]
    data_store: DataStore = request.app[app_keys.data_store]
    # find the requested stream
    if 'path' in request.query:
        stream = folder.find_stream_by_path(request.query['path'], db, stream_type=DataStream)
    elif 'id' in request.query:
        stream = db.get(DataStream,request.query["id"])
    else:
        return web.Response(text="specify an id or a path", status=400)
    if stream is None:
        return web.Response(text="stream does not exist", status=404)

    # parse time bounds
    start = None
    end = None
    try:
        if 'start' in request.query:
            start = int(request.query['start'])
        if 'end' in request.query:
            end = int(request.query['end'])
    except ValueError:
        return web.Response(text="[start] and [end] must be integers", status=400)

    # make sure bounds make sense
    if ((start is not None and end is not None) and
            (start >= end)):
        return web.Response(text="[start] must be < [end]", status=400)

    await data_store.remove(stream, start, end)
    return web.Response(text="ok")


async def consolidate(request):
    db: Session = request.app[app_keys.db]
    data_store: DataStore = request.app[app_keys.data_store]
    # find the requested stream
    if 'path' in request.query:
        stream = folder.find_stream_by_path(request.query['path'], db, stream_type=DataStream)
    elif 'id' in request.query:
        stream = db.get(DataStream,request.query["id"])
    else:
        return web.Response(text="specify an id or a path", status=400)
    if stream is None:
        return web.Response(text="stream does not exist", status=404)

    # parse time bounds
    start = None
    end = None
    try:
        if 'start' in request.query:
            start = int(request.query['start'])
        if 'end' in request.query:
            end = int(request.query['end'])
    except ValueError:
        return web.Response(text="[start] and [end] must be integers", status=400)

    # make sure bounds make sense
    if ((start is not None and end is not None) and
            (start >= end)):
        return web.Response(text="[start] must be < [end]", status=400)

    # parse the max_gap parameter
    if 'max_gap' not in request.query:
        return web.Response(text="specify max_gap as us integer", status=400)
    try:
        max_gap = int(request.query['max_gap'])
        if max_gap <= 0:
            raise ValueError()
    except ValueError:
        return web.Response(text="max_gap must be postive integer", status=400)
    num_removed = await data_store.consolidate(stream, start, end, max_gap)
    return web.json_response(data={"num_consolidated": num_removed})


async def decimate(request):
    db: Session = request.app[app_keys.db]
    data_store: DataStore = request.app[app_keys.data_store]
    # find the requested stream
    if 'path' in request.query:
        stream = folder.find_stream_by_path(request.query['path'], db, stream_type=DataStream)
    elif 'id' in request.query:
        stream = db.get(DataStream,request.query["id"])
    else:
        return web.Response(text="specify an id or a path", status=400)
    if stream is None:
        return web.Response(text="stream does not exist", status=404)
    await data_store.decimate(stream)
    return web.Response(text="ok")

async def drop_decimations(request):
    db: Session = request.app[app_keys.db]
    data_store: DataStore = request.app[app_keys.data_store]
    # find the requested stream
    if 'path' in request.query:
        stream = folder.find_stream_by_path(request.query['path'], db, stream_type=DataStream)
    elif 'id' in request.query:
        stream = db.get(DataStream, request.query["id"])
    else:
        return web.Response(text="specify an id or a path", status=400)
    if stream is None:
        return web.Response(text="stream does not exist", status=404)
    await data_store.drop_decimations(stream)
    return web.Response(text="ok")
