#
#   Imandra Inc.
#
#   codelogician/modeling/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,
    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 TaskKind(StrEnum):
    USER_IML_UPDATE = 'USER_IML_UPDATE'
    FORMALIZE = 'FORMALIZE'
    ANALYZE = 'ANALYZE'


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 CodeLogician to execute
    """

    # fmt: off
    task_kind        : TaskKind = TaskKind.FORMALIZE
    rel_path         : str
    src_code         : str = ''
    context          : str | None = ''
    dependencies     : list[FormalizationDependency] = Field(default_factory=list)
    graph_state      : GraphState | None = None
    language         : str = 'Python'
    mode             : ModelTaskMode = ModelTaskMode.AGENT
    gen_vgs          : bool = False
    gen_decomps      : bool = False
    # fmt: on

    # If set, we disregard everything else and just return these
    # along with the existing state.
    specified_commands: list[Command] | None = None

    task_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    created_at: datetime.datetime = Field(default_factory=datetime.datetime.now)

    def is_empty(self) -> bool:
        return self.src_code.strip() == ''

    def commands(self) -> list[Command]:
        """
        Return a list of CodeLogician agent commands for this task.
        """
        if self.specified_commands is not None:
            return self.specified_commands

        cmds: list[Command] = []

        # Always init from current src
        cmds.append(InitStateCommand(src_code=self.src_code, src_lang=self.language))

        # Optional context
        if self.context is not None:
            cmds.append(InjectFormalizationContextCommand(context=self.context))

        # Optional deps
        if self.dependencies:
            cmds.append(
                EditStateElementCommand(update={'dependency': self.dependencies})
            )

        # Mode
        if self.mode == ModelTaskMode.MANUAL:
            cmds.append(GenFormalizationDataCommand())
            cmds.append(GenModelCommand())
        elif self.mode == ModelTaskMode.AGENT:
            cmds.append(
                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}')

        # Optional artifacts
        if self.gen_vgs:
            cmds.append(GenVgsCommand(description=''))

        if self.gen_decomps:
            cmds.append(GenRegionDecompsCommand(function_name=None))

        return cmds

    def toJSON(self) -> str:
        return self.model_dump_json()

    @staticmethod
    def fromJSON(j: dict | str) -> 'ModelTask':
        if isinstance(j, str):
            return ModelTask.model_validate_json(j)
        return ModelTask.model_validate(j)

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


class DecompVGsModelTask(ModelTask):
    """
    Here we just perform extraction/running of Decomp/VGs
    """

    task_kind: TaskKind = TaskKind.ANALYZE

    def commands(self) -> list[Command]:
        if self.specified_commands is not None:
            return self.specified_commands

        if self.graph_state is None:
            raise ValueError(
                'DecompVGsModelTask requires graph_state; run formalization first.'
            )

        cmds: list[Command] = []
        if self.context is not None:
            cmds.append(InjectFormalizationContextCommand(context=self.context))
        if self.dependencies:
            cmds.append(
                EditStateElementCommand(update={'dependency': self.dependencies})
            )

        if self.gen_vgs:
            cmds.append(GenVgsCommand(description=''))
        if self.gen_decomps:
            cmds.append(GenRegionDecompsCommand(function_name=None))

        return cmds

    def __repr__(self) -> str:
        return f'DecompVGsModelTask with ID = {self.task_id}; path={self.rel_path}'


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

    task_kind: TaskKind = TaskKind.USER_IML_UPDATE

    iml_code: str  # user provided IML code

    def is_empty(self) -> bool:
        return False

    def commands(self) -> list[Command]:
        """
        We're just going to invoke the commands.
        """
        if self.graph_state:
            return [SetModelCommand(model=self.iml_code)]

        return [
            InitStateCommand(src_code=self.src_code, src_lang=self.language),
            SetModelCommand(model=self.iml_code),
        ]

    def __repr__(self) -> str:
        return f'UserManualIMLEditTask with ID = {self.task_id}; path={self.rel_path}'
