#
#   Imandra Inc.
#
#   codelogician/modeling/model_task_creator.py
#

from __future__ import annotations

import datetime
import logging
import re
from dataclasses import dataclass

from .dep_context import ModelDepContext
from .model import Model
from .status import ModelStatus
from .task import (
    DecompVGsModelTask,
    FormalizationNeed,
    ModelTask,
    TaskKind,
    UserManualIMLEditTask,
)

logger = logging.getLogger(__name__)


def _remove_import_directives(text: str) -> str:
    return re.sub(r'^\s*\[@@@import.*?\]\s*\n?', '', text, flags=re.MULTILINE)


@dataclass(frozen=True)
class ModelTaskCreator:
    user_wait_time: int | None = 5
    src_wait_time: int | None = 5

    # -------------------------
    # Checks / reasons
    # -------------------------

    def user_iml_change_ready(self, model: Model, wait: int | None = None) -> bool:
        if model.user_iml_edit is None:
            logger.debug('[%s] No user IML edit present', model.rel_path)
            return False

        wait_time = self.user_wait_time if wait is None else wait
        if wait_time is None:
            logger.debug(
                '[%s] User IML edit ready immediately (no wait)', model.rel_path
            )
            return True

        delta = datetime.datetime.now() - model.user_iml_edit.time_of_edit
        ready = delta > datetime.timedelta(seconds=wait_time)

        logger.debug(
            '[%s] User IML edit age=%ss wait=%ss ready=%s',
            model.rel_path,
            delta.total_seconds(),
            wait_time,
            ready,
        )
        return ready

    def has_src_code_changed(self, model: Model, wait: int | None = None) -> bool:
        if model.agent_state is None:
            logger.debug('[%s] No agent state → src considered changed', model.rel_path)
            return True

        wait_time = self.src_wait_time if wait is None else wait
        if wait_time is None or model.src_code_last_changed is None:
            long_enough = True
            delta_s: float | None = None
        else:
            delta = datetime.datetime.now() - model.src_code_last_changed
            delta_s = delta.total_seconds()
            long_enough = delta > datetime.timedelta(seconds=wait_time)

        changed = hash(model.src_code) != hash(model.agent_state.src_code)

        logger.debug(
            '[%s] Src change check: last_changed_age_s=%s wait=%s long_enough=%s changed=%s',
            model.rel_path,
            delta_s,
            wait_time,
            long_enough,
            changed,
        )
        return long_enough and changed

    def context_provided(self, model: Model) -> bool:
        provided = bool(model.context)
        logger.debug('[%s] Context provided=%s', model.rel_path, provided)
        return provided

    def formalization_reasons(
        self, model: Model, wait: int | None = None
    ) -> list[FormalizationNeed]:
        if model.src_code is None:
            logger.debug('[%s] No src code → no formalization reasons', model.rel_path)
            return []

        if model.agent_state is None:
            logger.debug('[%s] No agent state → needs formalization', model.rel_path)
            return [FormalizationNeed.NO_AGENT_STATE]

        reasons: list[FormalizationNeed] = []
        if self.has_src_code_changed(model, wait):
            reasons.append(FormalizationNeed.SRC_CODE_CHANGED)
        if self.context_provided(model):
            reasons.append(FormalizationNeed.CONTEXT_ADDED)
        if ModelDepContext.deps_changed(model):
            reasons.append(FormalizationNeed.DEPS_CHANGED)

        logger.debug('[%s] Formalization reasons=%s', model.rel_path, reasons)
        return reasons

    def deps_need_formalization(self, m: Model) -> bool:
        # fmt: off
        def p(d, t):
            r = self.formalization_reasons(d, t)
            if r:
                logger.debug('[%s] Dep %s needs formalization: %s', m.rel_path, d.rel_path, r)
            return r or self.deps_need_formalization(d)

        need = any(p(d, self.user_wait_time) for d in ModelDepContext.deps(m))
        # fmt: on
        logger.debug('[%s] deps_need_formalization=%s', m.rel_path, need)
        return need

    # -------------------------
    # Analysis (VGs / decomps)
    # -------------------------

    def _formalization_complete(self, model: Model) -> bool:
        """
        Conservative check: we consider the model "formalized enough to analyze"
        if we have an agent_state and some IML model exists.
        """
        complete = (model.agent_state is not None) and (model.iml_model() is not None)
        logger.debug('[%s] Formalization complete=%s', model.rel_path, complete)
        return complete

    def _missing_artifacts(self, model: Model) -> tuple[bool, bool]:
        """
        Returns: (need_vgs, need_decomps)

        Policy:
        - Empty results ([], None) can be a valid outcome of ANALYZE.
        - We only consider artifacts "missing" if we have not attempted generation
        since the current formalization baseline.
        """
        if model.agent_state is None:
            logger.debug(
                '[%s] Missing artifacts check skipped: no agent_state', model.rel_path
            )
            return (False, False)

        vgs = getattr(model.agent_state, 'vgs', None)
        decomps = getattr(model.agent_state, 'region_decomps', None)

        need_vgs = not model.analysis_vgs_attempted
        need_decomps = not model.analysis_decomps_attempted

        logger.debug(
            '[%s] Missing artifacts: need_vgs=%s need_decomps=%s '
            '(attempted_vgs=%s attempted_decomps=%s vgs_present=%s decomps_present=%s)',
            model.rel_path,
            need_vgs,
            need_decomps,
            model.analysis_vgs_attempted,
            model.analysis_decomps_attempted,
            bool(vgs),
            bool(decomps),
        )
        return (need_vgs, need_decomps)

    def _analysis_needed(self, model: Model) -> bool:
        """
        Analyze if formalization is complete, deps are stable, and artifacts are missing.
        """
        if not self._formalization_complete(model):
            logger.debug(
                '[%s] Analysis skipped: formalization incomplete', model.rel_path
            )
            return False

        # Don’t analyze if we need (re)formalization or deps are unstable.
        if self.formalization_reasons(model, self.user_wait_time):
            logger.debug('[%s] Analysis skipped: formalization pending', model.rel_path)
            return False
        if self.deps_need_formalization(model):
            logger.debug(
                '[%s] Analysis skipped: deps need formalization', model.rel_path
            )
            return False
        if ModelDepContext.deps_changed(model):
            logger.debug('[%s] Analysis skipped: deps changed', model.rel_path)
            return False

        need_vgs, need_decomps = self._missing_artifacts(model)
        needed = need_vgs or need_decomps
        logger.debug('[%s] Analysis needed=%s', model.rel_path, needed)
        return needed

    # -------------------------
    # Decide + create a single task
    # -------------------------

    def decide_task_kind(self, model: Model) -> TaskKind | None:
        """
        Decide the single highest-priority task for this model.

        Latest policy:

        - USER_IML_UPDATE can preempt an outstanding task if ready.

        - FORMALIZE can preempt an outstanding task if deps are stable AND there is
        a *real* new-change reason (e.g., SRC_CODE_CHANGED / CONTEXT_ADDED / DEPS_CHANGED),
        NOT merely "NO_AGENT_STATE" (which just means we're still waiting for the
        outstanding formalization result).

        - ANALYZE never preempts: only when there is no outstanding task AND artifacts are missing.
        """

        # 1) User IML override (allowed even if outstanding exists)
        if self.user_iml_change_ready(model):
            logger.debug(
                '[%s] Task decision=USER_IML_UPDATE (preempt=%s)',
                model.rel_path,
                model.outstanding_task_ID is not None,
            )
            return TaskKind.USER_IML_UPDATE

        # 2) Compute formalization needs
        reasons = self.formalization_reasons(model, self.user_wait_time)
        needs_formalization = bool(reasons)
        deps_stable = not self.deps_need_formalization(model)

        logger.debug(
            '[%s] Task decision precheck: outstanding=%s needs_formalization=%s deps_stable=%s reasons=%s',
            model.rel_path,
            model.outstanding_task_ID,
            needs_formalization,
            deps_stable,
            reasons,
        )

        # 3) If anything is outstanding:
        #    - allow FORMALIZE only if there's a real-change reason (not NO_AGENT_STATE)
        #    - never ANALYZE
        if model.outstanding_task_ID is not None:
            # If we're already running a FORMALIZE task, do not spam another one.
            if model.outstanding_task_kind is TaskKind.FORMALIZE:
                return None

            if needs_formalization and deps_stable:
                # "NO_AGENT_STATE" is not a preemption trigger: it just means we're still waiting
                # for the already-outstanding formalization task to produce an agent_state.
                real_change_reasons = [
                    r for r in reasons if r is not FormalizationNeed.NO_AGENT_STATE
                ]

                if real_change_reasons:
                    logger.debug(
                        '[%s] Task decision=FORMALIZE (preempt=True real_change_reasons=%s all_reasons=%s)',
                        model.rel_path,
                        real_change_reasons,
                        reasons,
                    )
                    return TaskKind.FORMALIZE

            logger.debug(
                '[%s] Task decision=NONE (blocked by outstanding_task_ID=%s; reasons=%s)',
                model.rel_path,
                model.outstanding_task_ID,
                reasons,
            )
            return None

        # 4) Normal formalize when idle
        if needs_formalization and deps_stable:
            logger.debug(
                '[%s] Task decision=FORMALIZE (preempt=False reasons=%s)',
                model.rel_path,
                reasons,
            )
            return TaskKind.FORMALIZE

        # 5) Analyze only when idle
        if self._analysis_needed(model):
            logger.debug('[%s] Task decision=ANALYZE', model.rel_path)
            return TaskKind.ANALYZE

        logger.debug('[%s] Task decision=NONE', model.rel_path)
        return None

    def create_task(self, model: Model) -> ModelTask | UserManualIMLEditTask | None:
        """
        Create at most one task.
        """

        if model.src_code is None:
            logger.debug(
                '[%s] Skipping task creation: src_code is None (deleted)',
                model.rel_path,
            )
            return None

        k = self.decide_task_kind(model)
        if k is None:
            logger.debug('[%s] No task created', model.rel_path)
            return None

        logger.debug('[%s] Creating task kind=%s', model.rel_path, k)

        if k is TaskKind.USER_IML_UPDATE:
            return self._user_iml_update_task(model)

        if k is TaskKind.FORMALIZE:
            return self._formalization_task(model)

        if k is TaskKind.ANALYZE:
            return self._analysis_task(model)

        raise RuntimeError(f'Unhandled task kind: {k}')

    # -------------------------
    # Task constructors
    # -------------------------

    def _formalization_task(self, model: Model) -> ModelTask:
        context = """Do not use `[@@measure ...]` annotations on nested functions (closures). If a function needs a `[@@measure ...]` annotation, structure the code so that the nested function is moved to the top-level.

When writing a recursive function where the recursion is controlled by an integer argument, put that integer argument as the first argument. Also try to arrange for the integer argument to count down to zero.

Use the operator `=` to test equality between numbers of type `Real`, not the operator `=.`, which does not exist."""

        task = ModelTask(
            rel_path=model.rel_path,
            src_code=model.src_code or '',
            context=f'{context}\n{model.context}',
            graph_state=model.agent_state.graph_state if model.agent_state else None,
            dependencies=ModelDepContext.dependencies_iml_models(model),
        )
        model.outstanding_task_ID = task.task_id
        logger.debug('[%s] Created FORMALIZE task_id=%s', model.rel_path, task.task_id)
        return task

    def _user_iml_update_task(self, model: Model) -> UserManualIMLEditTask | None:
        if model.user_iml_edit is None:
            logger.debug(
                '[%s] USER_IML_UPDATE requested but no user_iml_edit', model.rel_path
            )
            return None

        task = UserManualIMLEditTask(
            rel_path=model.rel_path,
            graph_state=model.agent_state.graph_state if model.agent_state else None,
            iml_code=_remove_import_directives(model.user_iml_edit.user_iml_entry),
            dependencies=ModelDepContext.dependencies_iml_models(model),
        )
        model.outstanding_task_ID = task.task_id
        logger.debug(
            '[%s] Created USER_IML_UPDATE task_id=%s', model.rel_path, task.task_id
        )
        return task

    def _analysis_task(self, model: Model) -> DecompVGsModelTask:
        need_vgs, need_decomps = self._missing_artifacts(model)

        task = DecompVGsModelTask(
            rel_path=model.rel_path,
            src_code=model.src_code or '',
            context=model.context,
            graph_state=model.agent_state.graph_state if model.agent_state else None,
            dependencies=ModelDepContext.dependencies_iml_models(model),
            gen_vgs=need_vgs,
            gen_decomps=need_decomps,
        )
        model.outstanding_task_ID = task.task_id
        logger.debug(
            '[%s] Created ANALYZE task_id=%s gen_vgs=%s gen_decomps=%s',
            model.rel_path,
            task.task_id,
            need_vgs,
            need_decomps,
        )
        return task

    # -------------------------
    # Status
    # -------------------------

    def gen_model_status(self, m: Model) -> ModelStatus:
        status = ModelStatus(
            src_code_status=m.src_code_status(),
            instructions_added=m.context_provided(),
            deps_changed=ModelDepContext.deps_changed(m),
            deps_need_formalization=self.deps_need_formalization(m),
            outstanding_task_ID=m.outstanding_task_ID,
            outstanding_task_kind=m.outstanding_task_kind,
            formalization_status=m.formalization_status(),
        )
        logger.debug('[%s] ModelStatus=%s', m.rel_path, status)
        return status
