#
#   Imandra Inc.
#
#   pyiml_strategy.py
#

from __future__ import annotations

import logging
from queue import Empty, Queue
from threading import Thread

from rich.progress import Progress, TaskID

from ..server.events import ServerEvent
from .config import StratConfig
from .events import (
    AutoModeCLTaskEvent,
    CLResultEvent,
    StrategyEvent,
)
from .pyiml_processor import PyIMLStrategyProcessor

log = logging.getLogger(__name__)


class PyIMLStrategy(Thread):
    """
    Strategy thread wrapper.

    The actual behavior is implemented in `PyIMLStrategyProcessor` for testability.
    """

    def __init__(
        self,
        state,
        config: StratConfig,
        oneshot: bool = False,
    ):
        super().__init__(daemon=True)

        # Processor owns all logic/stateful collaborators (worker, timers, etc.)
        self._proc = PyIMLStrategyProcessor(state=state, config=config, oneshot=oneshot)

        # This will be our events queue (we use `None` to stop the thread)
        self._queue: Queue[StrategyEvent | ServerEvent | None] = Queue()

        # Wire processor's callback to enqueue events
        self._proc.set_emit(self.add_event)

        log.info('Started PyIMLStrategy thread wrapper')

    # -------------------------
    # Passthroughs (public API)
    # -------------------------

    def language(self) -> str:
        return self._proc.language()

    def extensions(self) -> list[str]:
        return self._proc.extensions()

    def config(self):
        return self._proc.config

    def update_config(self, new_config) -> None:
        self._proc.update_config(new_config)

    def state(self):
        return self._proc.state

    def add_event(self, event: StrategyEvent | ServerEvent | None) -> None:
        """
        Add event to the queue for processing (None will force shutdown)
        """
        self._queue.put(event)
        log.info(
            f'Event added: {type(event).__name__ if event is not None else "None"}'
        )

    # -------------------------
    # Thread loop
    # -------------------------

    def run(self) -> None:
        # Let processor initialize (worker thread, initial sync, etc.)
        self._proc.start()

        if self._proc.oneshot:
            self._run_oneshot()
            return

        while True:
            try:
                event = self._queue.get()  # Blocks until an item is available

                if event is None:
                    log.info("Received a 'None' event to process. Shutting down.")
                    self._proc.shutdown()
                    break

                self._proc.process_event(event)
                self._queue.task_done()

            except Empty:
                pass
            except Exception as e:
                log.error(f'Caught exception during event processing: {str(e)}')

    def _run_oneshot(self) -> None:
        """
        One-shot loop. (Progress UI stays here; the business logic stays in processor.)
        """

        # We'll keep the progress of the various tasks here...
        progressTasks: dict[str, TaskID] = {}

        mmodel = self._proc.state.curr_meta_model
        if not mmodel:
            log.warning('Oneshot requested but no current MetaModel. Shutting down.')
            self._proc.shutdown()
            return

        with Progress() as progress:
            for path in mmodel.models:
                progressTasks[path] = progress.add_task(
                    f'[green] Formalizing {path}', total=1000
                )

            while True:
                try:
                    event = self._queue.get()

                    if event is None:
                        log.info("Received a 'None' event to process. Shutting down.")
                        self._proc.shutdown()
                        break

                    # Progress updates based on event type
                    if isinstance(event, AutoModeCLTaskEvent):
                        progress.update(
                            task_id=progressTasks[event.task.rel_path], completed=250
                        )
                    elif isinstance(event, CLResultEvent):
                        progress.update(
                            task_id=progressTasks[event.result.task.rel_path],
                            completed=1000,
                        )

                    self._proc.process_event(event)
                    self._queue.task_done()

                except Empty:
                    pass
                except Exception as e:
                    log.error(f'Caught exception during event processing: {str(e)}')
