"""Handles dynamic task outputs"""

import asyncio
import json
import logging
import queue
import threading
from collections import defaultdict
from functools import cached_property
from pathlib import Path
from typing import Callable, TYPE_CHECKING

from watchdog.events import FileSystemEventHandler

from experimaestro.ipc import ipcom
from experimaestro.utils import logger

from .base import Job, experiment

if TYPE_CHECKING:
    from experimaestro.core.objects import WatchedOutput


class TaskOutputCallbackHandler:
    def __init__(self, converter: Callable):
        pass


class TaskOutputs(FileSystemEventHandler):
    """Represent and monitors dynamic outputs generated by one task"""

    #: Global dictionary for handles
    HANDLERS: dict[Path, "TaskOutputs"] = {}

    #: Global lock to access current HANDLERS
    LOCK = threading.Lock()

    def create(job: Job):
        with TaskOutputs.LOCK:
            if instance := TaskOutputs.get(job.task_outputs_path, None):
                return instance

            instance = TaskOutputs(job.task_outputs_path)
            TaskOutputs[job.task_outputs_path] = instance
            return instance

    def __init__(self, path: Path):
        """Monitors an event path"""
        logger.debug("Watching dynamic task outputs in %s", path)
        self.path = path
        self.handle = None
        self.count = 0
        self.lock = threading.Lock()
        self.listeners: dict[str, dict[Callable, set[Callable]]] = defaultdict(
            lambda: defaultdict(set)
        )

        #: The events registered so far
        self.events = []

    def __enter__(self):
        """Starts monitoring task outputs"""
        self.job.task_outputs_path.parent.mkdir(parents=True, exist_ok=True)
        with self.lock:
            if self.handle is None:
                assert self.count == 0
                self.handle = ipcom().fswatch(self, self.path.parent, False)
            self.count += 1
        return self

    def __exit__(self, *args):
        """Stops monitoring task outputs"""
        with self.lock:
            self.count -= 1
            if self.count == 0:
                ipcom().fsunwatch(self.handle)
                self.fh.close()

                self.handle = None
                self._fh = None

    def watch_output(self, watched: "WatchedOutput"):
        """Add a new listener"""
        key = f"{watched.config.__identifier__}/{watched.method_name}"
        with self.lock:
            # Process events so far
            listener = self.listeners[key].get(watched.method, None)
            if listener is None:
                listener = TaskOutputCallbackHandler(watched.method)

            # Register
            self.listeners[key][watched.method].add(watched.callback)

    #
    # --- Events
    #

    @cached_property
    def fh(self):
        if self._fh is None:
            self._fh = self.path.open("rt")
        return self._fh

    def on_modified(self, event):
        self.handle(Path(event.src_path))

    def on_created(self, event):
        self.handle(Path(event.src_path))

    def handle(self, path: Path):
        if path != self.path:
            return

        with self.lock:
            logger.debug("[TASK OUTPUT] Handling task output for %s", self.path)

            while json_line := self.fh.readline():
                # Read the event
                event = json.loads(json_line)
                logger.debug("Event: %s", event)

                # FIXME: move elsewhere
                # # Process the event
                # event = self.config_method(
                #     self.job.config.__xpm__.mark_output,
                #     *event["args"],
                #     **event["kwargs"],
                # )

                self.events.append(event)
                # self.job.scheduler.xp.taskOutputsWorker.add(self, event)


class TaskOutputsWorker(threading.Thread):
    """This worker process dynamic output queue for one experiment"""

    def __init__(self, xp: experiment):
        super().__init__(name="task outputs worker")
        self.queue = queue.Queue()
        self.xp = xp

    def watch_output(self, watched: "WatchedOutput"):
        """Watch an output

        :param watched: The watched output specification
        """
        logger.debug("Registering task output listener %s", watched)

        # path = watched.job.tasks_output_path
        TaskOutputs.create(watched.job).watch_output(watched)

    def add(self, watcher, event):
        asyncio.run_coroutine_threadsafe(
            self.xp.update_task_output_count(1),
            self.xp.scheduler.loop,
        ).result()
        self.queue.put((watcher, event))

    def run(self):
        logging.debug("Starting output listener queue")
        while True:
            # Get the next element in the queue
            element = self.queue.get()
            if element is None:
                # end of processing
                break

            # Call all the listeners
            logging.debug("Got one event: %s", element)
            watcher, event = element
            for listener in watcher.listeners:
                try:
                    logger.debug("Calling listener [%s] with %s", listener, event)
                    listener(event)
                    logger.debug(
                        "[done] Calling listener [%s] with %s", listener, event
                    )
                except Exception:
                    logging.exception("Exception while calling the listener")
                self.queue.task_done()

                asyncio.run_coroutine_threadsafe(
                    self.xp.update_task_output_count(-1), self.xp.scheduler.loop
                ).result()
