#
#   Imandra Inc.
#
#   processor.py
#

from __future__ import annotations

import logging
import os
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from threading import Timer

from ..modeling.dep_context import ModelDepContext
from ..modeling.task import ModelTask, UserManualIMLEditTask
from ..server.events import FileSystemCheckbackEvent, FileSystemEvent, ServerEvent
from ..sketch.task import SketchChangeTask
from .config import StratConfig, StratMode
from .events import (
    AutoModeCLTaskEvent,
    ChangeMode,
    CLResultEvent,
    ModelCLTaskEvent,
    RunOneshotEvent,
    SketchChangeEvent,
    SketchChangeResultEvent,
    StrategyEvent,
)
from .state import StrategyState
from .worker import CodeLogicianWorker

log = logging.getLogger(__name__)

EmitFn = Callable[[StrategyEvent | ServerEvent | None], None]


@dataclass
class PyIMLStrategyProcessor:
    """
    Testable event processor for PyIMLStrategy.

    - No Queue management here.
    - Emits follow-up events via `self.emit(...)`.
    """

    state: StrategyState
    config: StratConfig
    oneshot: bool = False

    _emit: EmitFn | None = None
    _worker: CodeLogicianWorker | None = None

    _language: str = 'Python'
    _extensions: list[str] = None  # type: ignore[assignment]

    def __post_init__(self) -> None:
        if self._extensions is None:
            self._extensions = ['.py', '.iml']

        if self.config is None:
            self.config = StratConfig()

        # clear any task IDs that might have been persisted
        self.state.clear_task_ids()

    # -------------------------
    # Wiring
    # -------------------------

    def set_emit(self, emit: EmitFn) -> None:
        self._emit = emit

    def emit(self, event: StrategyEvent | ServerEvent | None) -> None:
        if self._emit is None:
            raise RuntimeError('Processor emit callback is not set')
        self._emit(event)

    # -------------------------
    # Public helpers
    # -------------------------

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

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

    def update_config(self, new_config: StratConfig) -> None:
        self.config = new_config

    # -------------------------
    # Lifecycle
    # -------------------------

    def start(self) -> None:
        """
        Initialize worker, do initial sync if needed.
        """
        if self._worker is None:
            self._worker = CodeLogicianWorker(self.emit)
            self._worker.start()

        if self.config.mode == StratMode.AUTO:
            log.info('Starting file sync')
            self.state.run_file_sync()
            self._check_autoformalization_tasks()

        log.info('Started CodeLogician Worker thread')

    def shutdown(self) -> None:
        """
        Shut down worker (by sending None task).
        """
        if self._worker is not None:
            self._worker.add_task(None)

    # -------------------------
    # Event dispatch
    # -------------------------

    def process_event(self, event: ServerEvent | StrategyEvent) -> None:
        if isinstance(event, FileSystemEvent):
            log.info('About to process file system update event')
            self._proc_filesystem_update(event)
            return

        if isinstance(event, FileSystemCheckbackEvent):
            log.info('About to process file system checkback event')
            self._proc_filesystem_checkback(event)
            return

        if isinstance(event, ChangeMode):
            log.info(f'About to process mode change to: {str(event)}')
            self._proc_mode_change(event.new_mode)
            return

        if isinstance(event, CLResultEvent):
            log.info('About to process CL Result event')
            self._on_cl_result(event)
            return

        if isinstance(event, RunOneshotEvent):
            log.info('About to process RunOneshotEvent')
            self._proc_oneshot_event(event)
            return

        if isinstance(event, ModelCLTaskEvent):
            log.info('About to process ModelCLTaskEvent')
            self._proc_model_cl_task_event(event)
            return

        if isinstance(event, AutoModeCLTaskEvent):
            log.info('About to process AutoModeCLTaskEvent')
            self._proc_auto_mode_cl_task_event(event)
            return

        if isinstance(event, SketchChangeEvent):
            self._proc_sketch_change_event(event)
            return

        if isinstance(event, SketchChangeResultEvent):
            self._proc_sketch_change_result(event)
            return

        log.error(f'Gotten an unknown event type: {type(event).__name__}')

    # -------------------------
    # Sketch handling
    # -------------------------

    def _proc_sketch_change_event(self, event: SketchChangeEvent) -> None:
        if self._worker is None:
            raise RuntimeError('Worker not started')
        self._worker.add_task(
            SketchChangeTask(
                sketch_id=event.sketch_id, change=event.change, iml_code=''
            )
        )

    def _proc_sketch_change_result(self, event: SketchChangeResultEvent) -> None:
        log.info(f'Received ImandraX result for sketch={event.sketch_id}')
        sketch = self.state.sketches.from_id(event.sketch_id)
        if sketch:
            sketch.apply_change_result(event.change_result)
            log.info('Successfully applied it now.')
        else:
            log.error(f'Unknown sketch with id={event.sketch_id}')

    # -------------------------
    # Auto mode tasks
    # -------------------------

    def _proc_auto_mode_cl_task_event(self, event: AutoModeCLTaskEvent) -> None:
        if self._worker is None:
            raise RuntimeError('Worker not started')
        self._worker.add_task(event.task)
        log.info('Adding new task')
        self.state.add_task(event.task)

    def _proc_model_cl_task_event(self, event: ModelCLTaskEvent) -> None:
        mmodel = self.state.curr_meta_model
        if not mmodel or event.rel_path not in mmodel.models:
            log.error(
                "Somethings wrong - either the MetaModel doesn't exist or the model is missing "
                f'[{event.rel_path}]'
            )
            return

        oldGraphState = None
        model = mmodel.models[event.rel_path]
        if model.agent_state:
            oldGraphState = model.agent_state.graph_state

        clTask = ModelTask(
            rel_path=event.rel_path,
            specified_commands=[event.cmd],
            graph_state=oldGraphState,
        )

        log.info('Adding new task')
        self.state.add_task(clTask)

    def _proc_oneshot_event(self, _event: RunOneshotEvent) -> None:
        self.oneshot = True
        self.config.mode = StratMode.AUTO

        log.info(
            "We're now in AUTO mode - will check what we should do with autoformalization..."
        )
        self._check_autoformalization_tasks()

    def _proc_mode_change(self, newMode: StratMode) -> None:
        self.config.mode = newMode
        log.info(f'Setting strategy mode to {newMode}')

        if newMode == StratMode.AUTO:
            log.info("Since we're in AUTO mode now, will check formalization tasks")
            self._check_autoformalization_tasks()

    def _check_autoformalization_tasks(self) -> None:
        log.info('Checking autoformalization tasks!')

        tasks = self.state.get_next_tasks()
        if tasks:
            log.info(f'Adding {len(tasks)} new tasks')
            for task in tasks:
                self.emit(AutoModeCLTaskEvent(task=task))
            return

        log.info('No new formalization tasks present!')

        # oneshot: exit when no pending tasks
        if self.oneshot and self.state.tasks == []:
            log.info('There are no tasks now. Quitting!')
            self.emit(None)

    # -------------------------
    # File system event handling
    # -------------------------

    def _check_back_callback(self, abs_path: str) -> None:
        log.info(f'Adding FileSystemCheckbackEvent for {abs_path}')
        self.emit(FileSystemCheckbackEvent(abs_path=abs_path))

    def _proc_filesystem_update(self, event: FileSystemEvent) -> None:
        # Apply the immediate point update (src/iml content, create/delete/move, etc.)
        self.state.curr_meta_model.process_filesystem_event(event)

        if self.config.mode != StratMode.AUTO:
            log.info("No checkback thread started as we're in STANDBY mode")
            return

        # Only schedule checkback for file types we actually use for tasks/deps.
        ext = Path(event.abs_path1).suffix
        if ext not in ('.py', '.iml'):
            log.info(f'Skipping checkback for unsupported extension: {ext}')
            return

        log.info(
            "Since we're in AUTO mode, scheduling a Timer to check back and "
            'see if the change is stable enough to formalize.'
        )

        timer = Timer(
            interval=self.config.time_for_filesync_checkback,
            function=self._check_back_callback,
            args=(event.abs_path1,),
        )
        timer.start()

    def _proc_filesystem_checkback(self, event: FileSystemCheckbackEvent) -> None:
        log.info('Processing FileSystemCheckbackEvent')

        if self.state.curr_meta_model is None:
            log.info('No current MetaModel. Aborting.')
            return

        ext = Path(event.abs_path).suffix

        if ext == '.py':
            if not self.state.curr_meta_model.abs_path_in_models(event.abs_path):
                log.info(f'Model for path: {event.abs_path} not found. Aborting.')
                return

        elif ext == '.iml':
            model_path = str(Path(event.abs_path).with_suffix('.py'))
            if not self.state.curr_meta_model.abs_path_in_models(model_path):
                log.info(
                    f'IML code for model with path: {model_path} not found. Aborting.'
                )
                return
        else:
            return

        # Always resync models & dependencies once the filesystem has settled
        # Note: here we'll also check for cycles in dependency graphs
        log.info('Running MetaModel filesystem sync after checkback')
        try:
            self.state.curr_meta_model.run_file_sync()
        except Exception as e:
            log.error(f'Failed to run MetaModel file sync: {e}')
            return

        if self.config.mode == StratMode.AUTO:
            self._check_autoformalization_tasks()

    def _write_model_code(self, rel_path: str) -> None:
        mmodel = self.state.curr_meta_model
        if not mmodel:
            log.error(
                "Somehow we're processing a CLResult event but we have no current MetaModel!"
            )
            return

        models = mmodel.models
        if rel_path not in models:
            log.info('Somehow the result is for a non-existent model')
            return

        model = models[rel_path]

        if not model.iml_code():
            return

        if self.config.write_models_in_same_dir:
            file_name = Path(rel_path).with_suffix('.iml')
            tgt_path = os.path.join(self.state.src_dir_abs_path, file_name)

            log.info(f'Will now write out IML model to: {tgt_path}')
            try:
                with open(tgt_path, 'w') as outfile:
                    import_statements = ModelDepContext.create_import_iml_statements(
                        model
                    )
                    full_model = f'{import_statements}\n\n{model.iml_code()}'

                    model.expected_iml_on_disk = full_model
                    print(full_model, file=outfile)
            except Exception as e:
                log.error(f'Failed to write out IML file: [{tgt_path}]: {e}')

            log.info(f'Finished writing {tgt_path}')
            return

        artifactDirPath = os.path.join(
            self.state.src_dir_abs_path, self.config.artifact_dir
        )
        if not os.path.isdir(artifactDirPath):
            log.warning(
                f"Artifact directory {artifactDirPath} doesn't exist. Trying to create it."
            )
            try:
                os.mkdir(artifactDirPath)
            except Exception as e:
                log.error(
                    f'Failed to create artifact directory [{artifactDirPath}]: error: {str(e)}. Aborting.'
                )
                return

        imlDirPath = os.path.join(artifactDirPath, self.config.iml_dir)
        if not os.path.isdir(imlDirPath):
            log.warning(
                f"IML model directory {imlDirPath} doesn't exist. Trying to create it."
            )
            try:
                os.mkdir(imlDirPath)
            except Exception as e:
                log.error(
                    f'Failed to create IML artifact directory {imlDirPath}: error: {str(e)}. Aborting.'
                )
                return

        filename = os.path.basename(model.rel_path).replace('.py', '.iml')
        imlFullpath = os.path.join(imlDirPath, filename)

        try:
            with open(imlFullpath, 'w') as outfile:
                print(model.iml_code(), file=outfile)
        except Exception as e:
            log.error(f'Failed to write IML file: [{imlFullpath}]: {str(e)}')

    def _on_cl_result(self, event: CLResultEvent) -> None:
        log.info(f'Applying CL Result to the model: {event}')

        self.state.apply_cl_result(event.result)
        self.state.remove_task(event.result.task)

        if self.config.write_consolidated:
            log.warning('Writing out a consolidated model is not supported yet!')

        is_not_user_iml_edit = not isinstance(event.result.task, UserManualIMLEditTask)
        if self.config.write_models and is_not_user_iml_edit:
            self._write_model_code(event.result.task.rel_path)

        if self.config.mode == StratMode.AUTO:
            self._check_autoformalization_tasks()
