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

import datetime
import logging
from typing import cast

from imandra.u.agents.code_logician.base import (
    VG,
    FormalizationDependency,
    FormalizationStatus,
    RegionDecomp,
    TopLevelDefinition,
)
from pydantic import BaseModel, Field, model_serializer

from ..strategy.cl_agent_state import CLAgentState
from ..util import filter as cl_filter
from .status import SrcCodeStatus
from .task import FormalizationNeed, TaskKind

log = logging.getLogger(__name__)


class UserIMLEdit(BaseModel):
    """
    We'll just use this to keep track of user edits
    """

    user_iml_entry: str
    src_code_at_the_time: str
    time_of_edit: datetime.datetime


class Model(BaseModel):
    """
    Container for a model - representing a single model along with its dependendencies,
    formalization state, user-overriden IML code, etc...

    These are the primary events that the model handles:
    - 1. Update to the source code (will trigger CL agent state)
    - 2. Update to the IML code (by the user)
    - 3. Applying the result of running CL agent
    - 4. User freezes/unfreezes their IML code (to prevent it from
        being overriden by CL results)
    - 5. User resets the changes that he made to the model (to the last CL model available)

    The metamodel updates the model when those events occur (listed above). It then
    uses the Model status function to figure out wether a new event should be written
    or an update to IML code on disk

    Three ways that a status is accessed:
    - formalization_reasons(user_wait_time : int) -- returns a list of formalization reasons
      (why a CL task should be generated). The list may be comprised of the following:
        - FormalizationNeed.NO_AGENT_STATE - no formalization state exists
        - FormalizationNeed.SRC_CODE_CHANGED - underlying source code changed
        - FormalizationNeed.CONTEXT_ADDED - human has provided context for the model
        - FormalizationNeed.DEPS_CHANGED - formalization state of one or more of the dependencies changed

        -- if the returned list is empty, then no further action is required

    - iml_on_fs_update_ready() - should we update the iml model that's written to disk
        - Description: this can be set when we have an update to the model

    """

    rel_path: str
    src_code: str | None = None
    src_code_last_changed: datetime.datetime | None = None
    src_language: str = 'Python'

    # This contains the formalization state -
    # note it is updated by both - running CL on the source code changes (e.g. python)
    # and on user-manual edits to IML (we call a separate step for that)
    agent_state: CLAgentState | None = None

    # Here we account for user's changes to the model IML code
    # and keep track of how long it's been since they edited it
    # if it's been long enough and it's different that what we have
    # then we'll kick off a command to make it part of the Graph State
    user_iml_edit: UserIMLEdit | None = (
        None  # This is human provided IML code and timestamp
    )

    # If the user IML code is frozen, then even if the
    # CL comes up with an update to the model, then it would still not
    # override the file on disk
    iml_code_frozen: bool = False

    # Here we keep a record of what should be on disk
    # Because we have two sources of IML models on disk (agent and user)
    # we keep this here to check for differences if there's an update (i.e.
    # if we write the model to disk after CL results come in, then we update this
    # and when the FS event comes in for IML change, we disregard it)
    expected_iml_on_disk: str | None = None

    # This is human-provided context that is sent to the CL agent
    context: str | None = None

    # This is used to keep track of the latest TaskID that was generated
    # We can only apply changes to the CL agent state if the taskID of the
    # result andf what we're waiting for matches - otherwise we discard the results
    outstanding_task_ID: str | None = None
    outstanding_task_kind: TaskKind | None = None

    # This helps know if we should kick-off VG/Decomp task after formalization
    analysis_vgs_attempted: bool = False
    analysis_decomps_attempted: bool = False

    # Persisted fields (loaded from JSON)
    str_dependencies: list[str] = Field(default_factory=list)
    str_rev_dependencies: list[str] = Field(default_factory=list)

    # In-memory links only
    dependencies: list['Model'] = Field(default_factory=list, exclude=True)
    rev_dependencies: list['Model'] = Field(default_factory=list, exclude=True)

    @model_serializer(mode='wrap')
    def _serialize_with_str_edges(self, handler):
        data = handler(self)

        data['str_dependencies'] = [m.rel_path for m in self.dependencies]
        data['str_rev_dependencies'] = [m.rel_path for m in self.rev_dependencies]

        return data

    # This gets updated along with each agent state - here we maintain
    # IML code for the dependencies that were used. If they differ
    # then we may need to re-run CodeLogician.
    formalized_deps: list[FormalizationDependency] = Field(default_factory=list)

    # We do the same with the human-provided context
    formalized_context: str | None = None

    # Static formalization reasons - note if the model is "live"
    # (i.e. not post deserialization (so we have actual dependencies present)
    # then this is not used - this is only to store those for serialization
    # when we have to present the model
    static_frm_reasons: list[FormalizationNeed] = Field(default_factory=list)

    def add_dependency(self, d: 'Model') -> None:
        """
        "Safe" add of a dependency - if it's already there, it would not add it twice
        """
        if d not in cast(list[Model], self.dependencies):
            cast(list[Model], self.dependencies).append(d)

    def set_iml_code_frozen(self) -> None:
        """
        Set IML code frozen - this will prevent any
        """
        self.iml_code_frozen = True

    def set_iml_code_unfrozen(self) -> None:
        """
        Here we set the IML code unfrozen
        """
        self.iml_code_frozen = False

    def cl_result_available(self) -> bool:
        """
        Is there a CL result that has been available, but not set
        b/c of the frozen client model?
        """
        return False

    def context_provided(self) -> bool:
        """
        Return True if there's human-provided context
        """
        return bool(self.context)

    def iml_on_fs_update_ready(self) -> bool:
        """
        Returns True/False if there's an update to the IML model that should be written to disk...
        """
        return False

    def apply_agent_state(
        self,
        agent_state: CLAgentState,
        dependencies: list[FormalizationDependency] = [],
    ) -> None:
        """
        Apply the agent state and store the dependencies that were used in the call to CL agent

        Parameters:
        - agent_state - CLAgentState that we'll apply
        - dependencies - the dependencies that were used during formalization - we'll use them later to compare if
            we need to redo autoformalization.
        """

        # Let's also now reset the task ID
        self.outstanding_task_ID = None
        self.outstanding_task_kind = None

        # Agent state
        self.agent_state = agent_state

        # Remove instructions that were used from the current list of outstanding ones
        self.formalized_context = agent_state.context

        # Let's now set the dependencies' models that were used
        self.formalized_deps = dependencies

    def user_iml_change_ready(self, user_wait_time: int | None) -> bool:
        """
        Is the user's IML change ready to be turned into a task. Here, we consider whether
        the change done by the user is "old enough". We do this to ensure we're only formalizing
        changes that have become 'stable'.
        """

        if self.user_iml_edit is None:
            return False

        if user_wait_time is None:
            long_enough = True
        else:
            long_enough = (
                datetime.datetime.now() - self.user_iml_edit.time_of_edit
            ) > datetime.timedelta(seconds=user_wait_time)

        return long_enough

    def src_code_status(self, srcCodeWaitTime: int | None = None) -> SrcCodeStatus:
        """
        Return source code status
        """
        if self.src_code is None:
            return SrcCodeStatus.SRC_CODE_DELETED

        else:
            if srcCodeWaitTime is None or self.src_code_last_changed is None:
                timeLongEnough = True
            else:
                timeLongEnough = (
                    datetime.datetime.now() - self.src_code_last_changed
                ) > datetime.timedelta(seconds=srcCodeWaitTime)

        if timeLongEnough and (hash(self.src_code) == hash(self.iml_code())):
            return SrcCodeStatus.SRC_CODE_CHANGED
        else:
            return SrcCodeStatus.SRC_CODE_CURRENT

    def formalization_status(self) -> FormalizationStatus:
        """
        Return formalization status of the model
        """
        if self.agent_state is None:
            return FormalizationStatus.UNKNOWN
        else:
            return self.agent_state.status

    def has_src_code_changed(self, user_wait_time: int | None = None) -> bool:
        """
        Return True if source code has changed from when it was formalized
        """

        # if we're missing formalization altogether, then source code has changed
        if self.agent_state is None:
            return True

        if user_wait_time is None or self.src_code_last_changed is None:
            timeLongEnough = True
        else:
            timeLongEnough = (
                datetime.datetime.now() - self.src_code_last_changed
            ) > datetime.timedelta(seconds=user_wait_time)

        return timeLongEnough and (
            hash(self.src_code) != hash(self.agent_state.src_code)
        )

    def set_iml_model(self, new_iml_model: str, record_time: bool = False) -> None:
        """
        Set the IML model code and record the time if we need to. This is called
        when the user manually overrides the model. We also need to check if the model
        actually changed... (this event is all also  called when we write out the model
        from CL result - it can't tell the difference, so we need to check it here...)
        """

        expected = self.expected_iml_on_disk
        if expected is not None and new_iml_model.strip() == str(expected).strip():
            log.warning(
                f"Model set by the user doesn't appear to be any different from existing model: [{self.rel_path}]"
            )
            return

        log.info(f'User is updating the IML model: [{self.rel_path}]')

        self.user_iml_edit = UserIMLEdit(
            user_iml_entry=new_iml_model,
            src_code_at_the_time='N\\' if self.src_code is None else self.src_code,
            time_of_edit=datetime.datetime.now(),
        )

    def set_src_code(self, new_src_code: str | None, record_time: bool = False) -> None:
        """
        Set the src_code and update the time source code changed
        """

        # This means the code was deleted
        if new_src_code is None:
            self.src_code = None
            self.src_code_last_changed = None
            return

        # if not isinstance(new_src_code, str):
        #    raise Exception (f"Expected a string value for `new_src_code`, but got {type(new_src_code).__name__}")

        if hash(new_src_code) != hash(self.src_code):
            log.info(f'Updating source code for model [{self.rel_path}]')

            if record_time:
                self.src_code_last_changed = datetime.datetime.now()
            else:
                self.src_code_last_changed = None

            self.src_code = new_src_code
        else:
            log.warning(f"Specified model hasn't changed: [{self.rel_path}]")

    def iml_code(self) -> str | None:
        """
        If available, returns IML code with artifacts (e.g. VGs and decomps)
        """
        return self.agent_state.iml_code if self.agent_state else None

    def iml_model(self) -> str | None:
        """
        If available, returns IML model (not artifacts like VGs, decomps). This
        also looks at the available user-overriden code.
        """

        return self.agent_state.iml_model if self.agent_state else None

    def verification_goals(self) -> list[VG]:
        """
        If available, returns the list of dictionaries with verification goals
        """
        return self.agent_state.vgs if self.agent_state else []

    def failed_vgs(self) -> list[VG]:
        """
        Return just the list of failed VGs
        """
        return cl_filter(lambda x: not x['status'], self.verification_goals())

    def decomps(self) -> list[RegionDecomp]:
        """
        If available, returns the list of decomposition requests
        """
        return self.agent_state.region_decomps if self.agent_state else []

    def opaque_funcs(self) -> list[TopLevelDefinition]:
        """
        If available, return the opaque functions used along with their approximations
        """
        return self.agent_state.opaque_funcs if self.agent_state else []

    def gen_stats(self) -> dict[str, str | int]:
        """
        Generate some numberical stats for this model
        """
        s: dict[str, str | int] = {}
        s['frm_status'] = str(self.formalization_status())
        s['num_opaques'] = len(self.opaque_funcs())
        s['num_failed_vgs'] = len(self.failed_vgs())

        return s

    @staticmethod
    def fromJSON(j: str | dict[str, str]) -> 'Model':
        """
        Return a Model object from provided JSON. Note that we do not set the lists
        of models for 'dependsOn' here because we don't have access to the
        actual model objects. This is done outside this state methods.
        """

        if isinstance(j, str):
            return Model.model_validate_json(j)
        else:
            return Model.model_validate(j)

    def __hash__(self):
        """
        We need this so we can compare models
        """
        return hash(str(self.model_dump_json()))
