#
#   Imandra Inc.
#
#   model_task.py
#

import datetime
import uuid
from enum import StrEnum

from imandra.u.agents.code_logician.base import FormalizationDependency
from imandra.u.agents.code_logician.command import (
    AgentFormalizerCommand,
    Command,
    EditStateElementCommand,
    EmbedCommand,
    GenFormalizationDataCommand,
    GenModelCommand,
    GenRegionDecompsCommand,
    GenVgsCommand,
    InitStateCommand,
    InjectFormalizationContextCommand,
    SetModelCommand,
)
from imandra.u.agents.code_logician.graph import GraphState
from pydantic import BaseModel, Field


class ModelTaskMode(StrEnum):
    """
    Strategy mode
    """

    AGENT = 'Agent'
    MANUAL = 'Manual'


class FormalizationNeed(StrEnum):
    """
    Reasons to perform formalization
    """

    # fmt: off
    NO_AGENT_STATE = 'No_agent_state' # This is the first time we're formalizing the model
    SRC_CODE_CHANGED = 'Src_code_changed' # Source code has changed
    IML_CODE_CHANGED = 'IML_code_changed' # User edited the IML code, so we need to resend it to CL
    CONTEXT_ADDED = 'Cntxt_added' # Human feedback has been provided
    DEPS_CHANGED = 'Deps_changed' # Dependencies (at least one) have changed
    # fmt: on


class ModelTask(BaseModel):
    """
    Model task contains a single task for Code Logician to execute
    """

    # fmt: off
    rel_path        : str # path of the source file (relative to the source code directory)
    src_code        : str = '' # source code
    context         : str | None = '' # IML source code for dependent models
    dependencies    : list[FormalizationDependency] = [] # IML source code for dependent models
    graph_state     : GraphState | None = None # previous graph state, if available
    language        : str = 'Python' # programming language
    mode            : ModelTaskMode = ModelTaskMode.AGENT # mode used with CodeLogician (either 'simple' or 'agent')
    gen_vgs         : bool = False # Should we generate verification goals?
    gen_decomps     : bool = False # Should we generate decompositions?
    gen_embeddings  : bool = True # Should we generate embeddings?
    # fmt: on

    # This is used to submit user-specified commands for specific models
    # if this is set, then we disregard everything else and just return these
    # along with the existing state
    specified_commands: list[Command] | None = None

    # Each task has a unique ID that we'll then use to assign back the result
    # of CL work for this task - if the model's task ID doesn't match the result
    # then we discard it - this way we always keep the most recent result
    task_id: str = Field(
        default_factory=lambda: str(uuid.uuid4())
    )  # assigned task ID, new one created if not provided

    # Same with the timestamp, if it's provided (mostly during de/serialization,
    # then let's use it otherwise, let's create a new one)
    created_at: datetime.datetime = Field(
        default_factory=datetime.datetime.now
    )  # Each instance gets a new timestamp

    def commands(self):
        """
        Return a list of CodeLogician agent commands for this task
        """

        def yield_commands():
            yield InitStateCommand(src_code=self.src_code, src_lang=self.language)

            if self.context is not None:
                yield InjectFormalizationContextCommand(context=self.context)

            if self.dependencies:
                yield EditStateElementCommand(update={'dependency': self.dependencies})

            if self.mode == ModelTaskMode.MANUAL:
                yield GenFormalizationDataCommand()
                yield GenModelCommand()

            elif self.mode == ModelTaskMode.AGENT:
                yield AgentFormalizerCommand(
                    no_gen_model_hitl=True,
                    max_tries_wo_hitl=3,
                    max_tries=3,
                    no_check_formalization_hitl=True,
                    no_refactor=False,
                )
            else:
                raise Exception(f'Attempting unrecognized mode: {self.mode}')

            if self.gen_vgs:
                yield GenVgsCommand(description='')

            if self.gen_decomps:
                yield GenRegionDecompsCommand(function_name=None)
                # for i in self.tests:
                #     yield GenTestCasesCommand(decomp_idx=i)

            if self.gen_embeddings:
                yield EmbedCommand()

        return self.specified_commands or list(yield_commands())

    def toJSON(self):
        """
        Convert to a JSON
        """
        return self.model_dump_json()

    @staticmethod
    def fromJSON(j: dict | str):
        """
        fromJSON
        """
        if isinstance(j, str):
            return ModelTask.model_validate_json(j)
        else:
            return ModelTask.model_validate(j)

    def __repr__(self):
        """ """
        return f'Base ModelTask with ID = {self.task_id}; path={self.rel_path}'


class UserManualIMLEditTask(ModelTask, BaseModel):
    """
    This updates the graph state from a manual user IML edit
    """

    iml_code: str  # user provided IML code

    def commands(self):
        """
        We're just going to invoke the commands
        """

        if self.graph_state:
            commands = [SetModelCommand(model=self.iml_code)]
        else:
            commands: list[Command] = [
                InitStateCommand(src_code=self.src_code, src_lang=self.language),
                SetModelCommand(model=self.iml_code),
            ]

        return commands
