#  Copyright (c) 2020 KTH Royal Institute of Technology
#
#  Licensed 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.
#
#  Licensed 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.
from __future__ import annotations

import abc
from datetime import datetime
import threading
from collections import namedtuple
from pathlib import Path
from queue import Empty, Queue
from typing import Any, Mapping, NamedTuple, Sequence, Set, Union

import numpy as np
import pandas as pd

from .logging import Logger


class Recorder(abc.ABC):
    """
    Interface for Recorder objects. Recorder and Recordable objects together
    implement the Observable pattern for text records of the plant state.
    """

    def __init__(self, recordable: Recordable):
        recordable.recorders.add(self)

    @abc.abstractmethod
    def notify(self, latest_record: NamedTuple):
        """
        Notify this recorder of the latest record generated by a recordable
        it is monitoring.
        """
        pass

    def initialize(self):
        """
        Called on plant startup.
        """
        pass

    def shutdown(self):
        """
        Called on plant shutdown.
        """
        pass

    def flush(self):
        """
        Mainly intended for IO-backed recorders, to force flushing of any
        potentially memory-stored data.
        """
        pass


class Recordable(abc.ABC):
    """
    Interface for Recordable objects. Recorder and Recordable objects together
    implement the Observable pattern for text records of the plant state.
    """

    @property
    @abc.abstractmethod
    def recorders(self) -> Set[Recorder]:
        pass

    @property
    @abc.abstractmethod
    def record_fields(self) -> Sequence[str]:
        pass


class NamedRecordable(Recordable):
    """
    A base implementation of a named recordable object.
    """

    def __init__(self,
                 name: str,
                 record_fields: Sequence[str],
                 opt_record_fields: Mapping[str, Any] = {}):
        self._name = name
        all_record_fields = list(record_fields) + list(opt_record_fields.keys())
        self._record_cls = namedtuple('_Record', all_record_fields,
                                      defaults=opt_record_fields.values())
        self._record_fields = self._record_cls._fields
        self._recorders: Set[Recorder] = set()

    @property
    def record_fields(self) -> Sequence[str]:
        return self._record_fields

    def push_record(self, **kwargs) -> None:
        record = self._record_cls(**kwargs)
        for recorder in self._recorders:
            recorder.notify(record)

    @property
    def recorders(self) -> Set[Recorder]:
        return self._recorders


class CSVRecorder(Recorder):
    """
    Implementation of a CSV-file backed Recorder.
    """

    def __init__(self,
                 recordable: Recordable,
                 output_dir: Path,
                 metric_name: str,
                 chunk_size: int = 1000):
        # TODO: maybe add timestamp??
        super(CSVRecorder, self).__init__(recordable)

        self._recordable = recordable
        self._log = Logger()

        output_dir = output_dir.resolve()
        if output_dir.exists() and not output_dir.is_dir():
            raise FileExistsError(f'{self._path} exists and is not a '
                                  f'directory!')

        output_dir.mkdir(exist_ok=True, parents=True)
        self._path = output_dir / \
                     f'{metric_name}.{datetime.now():%Y%m%d.%H%M%S%f}.csv'

        dummy_data = np.empty((chunk_size, len(recordable.record_fields)))
        self._table_chunk = pd.DataFrame(data=dummy_data,
                                         columns=recordable.record_fields)
        self._chunk_count = 0
        self._chunk_row_idx = 0

        self._chunk_q = Queue()
        self._chunk_write_thread = threading.Thread(
            target=self._writer_loop, args=(self._path,))
        self._shutdown_event = threading.Event()

    def _writer_loop(self, path: Path) -> None:
        with path.open('wb') as fp:
            # "touch" the file to clear it and prepare for actual writing
            fp.write(bytes(0x00))

        with path.open('a', newline='') as fp:
            while not self._shutdown_event.is_set():
                try:
                    chunk, with_hdr = self._chunk_q.get(block=True,
                                                        timeout=0.05)
                except Empty:
                    continue

                chunk.to_csv(fp, header=with_hdr, index=False)
                self._chunk_q.task_done()

            # this only occurs at shutdown, so there shouldn't be any new
            # chunks coming in
            while not self._chunk_q.empty():
                chunk, with_hdr = self._chunk_q.get_nowait()
                chunk.to_csv(fp, header=with_hdr, index=False)
                self._chunk_q.task_done()

    def initialize(self) -> None:
        self._log.info(f'Initializing CSVRecorder on {self._path}.')
        # initialize the writing thread
        self._shutdown_event.clear()
        while self._chunk_q.qsize() > 0:
            self._chunk_q.get_nowait()
        self._chunk_write_thread.start()

    def flush(self) -> None:
        """
        Flushes the internal record buffer to the backing CSV file.
        """

        self._chunk_q.put((
            (self._table_chunk.iloc[:self._chunk_row_idx].copy()),
            self._chunk_count == 0))

        self._chunk_row_idx = 0
        self._chunk_count += 1

        return

    def notify(self, latest_record: NamedTuple) -> None:
        self._table_chunk.iloc[self._chunk_row_idx] = \
            pd.Series(latest_record._asdict())
        self._chunk_row_idx += 1

        if self._chunk_row_idx == self._table_chunk.shape[0]:
            # flush to disk
            self.flush()

    def shutdown(self) -> None:
        self._log.info(f'Flushing and closing CSV table writer on path '
                       f'{self._path}...')
        self.flush()
        self._shutdown_event.set()
        self._chunk_q.join()
        self._chunk_write_thread.join()
        self._chunk_write_thread = threading.Thread(
            target=self._writer_loop, args=(self._path,))
