#
#   Imandra Inc.
#
#   speclogician/modeling/domain.py
#

from typing import Sequence
from pydantic import BaseModel, Field

from .predicates import (
    StatePredicate, ActionPredicate, Transition, PredicateType
)
from .__init__ import extract_declared_types
from enum import StrEnum

class BaseStatus(StrEnum):
    UNKNOWN = "unknown"
    INVALID_IML = "invalid_iml"
    VALID = "valid"

class DomainModel(BaseModel):
    """ 
    The domain model (base model with state/action types + other auxillary functions,
    and predicates/transition functions)
    """

    base_status : BaseStatus = BaseStatus.UNKNOWN # IML status
    base_has_state : bool = False # Model contains `state` type definition
    base_has_action : bool = False # Model contains `action` type definition
    base : str = "" # Includes type definitions and other 

    transitions : list[Transition] = Field(default_factory=list[Transition]) # 'state * action -> state' tansition functions
    state_preds : list[StatePredicate] = Field(default_factory=list[StatePredicate]) # 'state -> bool' predicates
    action_preds : list[ActionPredicate] = Field(default_factory=list[ActionPredicate]) # 'state * action -> bool' predicates

    # transitions
    num_trans_total : int = 0
    num_trans_matched : int = 0
    num_trans_valid_logic : int = 0
    num_trans_errored : int = 0

    # state-only predicates
    num_state_preds_total : int = 0
    num_state_preds_matched : int = 0
    num_state_preds_valid_logic : int = 0
    num_state_preds_errored : int = 0

    # state x action predicates
    num_action_preds_total : int = 0
    num_action_preds_matched : int = 0
    num_action_preds_valid_logic : int = 0
    num_action_preds_errored : int = 0

    # predicate totals
    num_preds_total : int = 0
    num_preds_matched : int = 0
    num_preds_valid_logic : int = 0
    num_preds_errored : int = 0

    def assert_has_predicates(
        self,
        *,
        given: Sequence[str],
        when: Sequence[str],
    ) -> None:
        """Ensure all predicates referenced by a scenario exist in the domain model.

        - `given` must reference state predicates (state -> bool)
        - `when` must reference action predicates (state -> action -> bool)

        Raises ValueError with a helpful message if anything is missing or mismatched.
        """
        missing_state: list[str] = []
        missing_action: list[str] = []
        wrong_kind_given: list[str] = []
        wrong_kind_when: list[str] = []

        for name in given:
            p = self.get_pred_by_name(name)
            if p is None:
                missing_state.append(name)
            elif p not in self.state_preds:
                # Exists, but it's an action predicate (or otherwise not a state pred)
                wrong_kind_given.append(name)

        for name in when:
            p = self.get_pred_by_name(name)
            if p is None:
                missing_action.append(name)
            elif p not in self.action_preds:
                # Exists, but it's a state predicate (or otherwise not an action pred)
                wrong_kind_when.append(name)

        if missing_state or missing_action or wrong_kind_given or wrong_kind_when:
            parts: list[str] = []
            if missing_state:
                parts.append(f"missing state predicates (given): {missing_state}")
            if wrong_kind_given:
                parts.append(f"not state predicates (given): {wrong_kind_given}")
            if missing_action:
                parts.append(f"missing action predicates (when): {missing_action}")
            if wrong_kind_when:
                parts.append(f"not action predicates (when): {wrong_kind_when}")

            raise ValueError("Scenario references invalid predicates: " + "; ".join(parts))

    def assert_has_transitions(self, *, then: Sequence[str]) -> None:
        """Ensure all transitions referenced by a scenario exist in the domain model.

        `then` is a list of transition function names.

        Raises ValueError with a helpful message if anything is missing.
        """
        missing: list[str] = []

        for name in then:
            t = self.get_trans_by_name(name)
            if t is None:
                missing.append(name)

        if missing:
            raise ValueError(f"Scenario references unknown transitions (then): {missing}")

    def get_pred_by_name (self, name:str) -> StatePredicate|ActionPredicate|None:
        """ Return a predicate body for the specified name """
        s_pred = next((p for p in self.state_preds if p.name == name), None)

        if s_pred is None:
            return next((p for p in self.action_preds if p.name == name), None)
        else:
            return s_pred

    def get_trans_by_name (self, name:str) -> Transition|None:
        """ Return transition body for a specified name """
        return next((p for p in self.transitions if p.name == name), None)

    def action_pred_names(self) -> list[str]:
        """ List of action predicate names """
        return list(map(lambda x: x.name, self.action_preds))
    
    def state_pred_names(self) -> list[str]:
        """ List of available state predicate names """
        return list(map(lambda x: x.name, self.state_preds))

    def transition_names(self) -> list[str]:
        """ List of available transitions """
        return list(map(lambda x: x.name, self.transitions))

    def pred_exists(self, name:str) -> bool:
        """
        Does this predicate exist?
        """
        if next ((p for p in self.state_preds if p.name == name), None) is not None:
            return True    
        return next ((p for p in self.action_preds if p.name == name), None) is not None
    
    def add_predicate(self, name: str, ptype: PredicateType, src_code: str) -> None:
        """Add a new predicate to the domain model."""

        # Enforce global uniqueness across all predicates
        if any(p.name == name for p in self.state_preds) or \
        any(p.name == name for p in self.action_preds):
            raise ValueError(f"Predicate '{name}' already exists in domain model")

        if ptype == PredicateType.STATE:
            pred = StatePredicate(
                name=name,
                src_code=src_code,
            )
            self.state_preds.append(pred)

        elif ptype == PredicateType.ACTION:
            pred = ActionPredicate(
                name=name,
                src_code=src_code,
            )
            self.action_preds.append(pred)

        else:
            raise ValueError(f"Attempting to add unknown type of predicate: {ptype}")

    def edit_predicate(self, name: str, src_code: str) -> None:
        """Edit a predicate by name and update last_updated."""

        matches = [p for p in self.state_preds if p.name == name] + \
                [p for p in self.action_preds if p.name == name]

        if not matches:
            raise ValueError(f"Predicate '{name}' not found")

        if len(matches) > 1:
            raise ValueError(f"Predicate '{name}' is ambiguous (found {len(matches)} matches)")

        pred = matches[0]
        pred.set_src_code(src_code)

    def rem_predicate(self, name: str) -> None:
        """Remove a predicate (state or action) by name."""

        removed = False

        # Remove from state predicates
        for i, p in enumerate(self.state_preds):
            if p.name == name:
                del self.state_preds[i]
                removed = True
                break

        # Remove from action predicates
        for i, p in enumerate(self.action_preds):
            if p.name == name:
                del self.action_preds[i]
                removed = True
                break

        if not removed:
            raise ValueError(f"Predicate '{name}' not found in domain model")
        
    def add_transition(self, name: str, src_code: str) -> None:
        """Add a new transition function to the domain model."""

        # Ensure uniqueness
        if any(t.name == name for t in self.transitions):
            raise ValueError(f"Transition '{name}' already exists in domain model")

        transition = Transition(
            name=name,
            src_code=src_code,
        )

        self.transitions.append(transition)

    def edit_transition(self, name: str, src_code: str) -> None:
        """Edit a transition function by name and update last_updated."""
        matches = [t for t in self.transitions if t.name == name]

        if not matches:
            raise ValueError(f"Transition '{name}' not found in domain model")

        if len(matches) > 1:
            raise ValueError(
                f"Transition name '{name}' is ambiguous "
                f"({len(matches)} transitions found)"
            )

        transition = matches[0]
        transition.set_src_code(src_code)

    def rem_transition(self, name: str) -> None:
        """Remove a transition function by name."""

        for i, t in enumerate(self.transitions):
            if t.name == name:
                del self.transitions[i]
                return

        raise ValueError(f"Transition '{name}' not found in domain model")

    def trans_exists(self, name:str) -> bool:
        """ Return True/False if transitions with specified name exists """
        return next ((t for t in self.transitions if t.name == name), None) is not None

    def state_specified (self) -> bool:
        """
        Return True if the `state` type is properly specified
        """
        return self.base_has_type('state')
    
    def action_specified(self) -> bool:
        """
        Return True if the `action` is properly specified
        """
        return self.base_has_type('action')

    def base_has_type(
            self, 
            type_name : str
        ) -> bool:
        """
        Check whether the base model has a specified type
        """

        declared_types = extract_declared_types(self.base)
        return type_name in declared_types
    
    def set_base (
            self, 
            new_base : str) -> None:
        """
        Update the base model
        """
        self.base = new_base

        return None

    def make_iml_model (
        self, 
        state_pred_names : list[str], 
        action_pred_names : list[str],
        transition_names : list[str]
        ) -> str:
        """ 
        Return base model with required predicates/transitions if needed 
        """

        state_preds_str = ""
        for p in state_pred_names:
            pred = self.get_pred_by_name(p)
            if pred is None:
                continue
            state_preds_str += f"{pred.to_iml()}"
        if len(state_preds_str):
            state_preds_str = f"(* State predicates *)\n{state_preds_str}"
        
        action_preds_str = ""
        for p in action_pred_names:
            pred = self.get_pred_by_name(p)
            if pred is None:
                continue
            action_preds_str += f"\n{pred.to_iml()}\n"
        if len(action_preds_str):
            action_preds_str = f"(* Action predicates *)\n{action_preds_str}"

        trans_str = ""
        for t in transition_names:
            trans = self.get_trans_by_name(t)
            if trans is None:
                continue
            trans_str += f"\n{trans.to_iml()}\n"
        if len(trans_str):
            trans_str = f"(* Transitions *)\n{trans_str}"

        return f"""
(* Domain Model *)

(* Base *)
{self.base}
{state_preds_str}
{action_preds_str}
{trans_str}
""".rstrip("\n")

    def to_iml(self):
        """
        Create a consolidated IML model with all predicates and transitions
        """
        return self.make_iml_model(
            self.state_pred_names(), self.action_pred_names(), self.transition_names()
        )
