#
#   Imandra Inc.
#
#   speclogician/state/state.py
#

import os
from pathlib import Path
from pydantic import BaseModel, Field

from ..modeling.spec_stats import refresh_spec_stats
from ..modeling.trace_reachability_compute import compute_trace_reachability_analysis
from ..logic import analyze_change
from ..data.refs import DocRef, SrcCodeRef
from ..data.traces import TestTrace, LogTrace
from ..utils import console
from .inst import StateInstance, InstancesList
from .diff_calcs import calc_diff
from .recommender import Recommender, Recommendation, ChangeFailure

from .change import (
    StateChange,
    DomainModelBaseEdit,
    PredicateAdd,
    PredicateEdit,
    PredicateRemove,
    TransitionAdd,
    TransitionEdit,
    TransitionRemove,
    ScenarioAdd,
    ScenarioEdit,
    ScenarioRemove,
    AddTestTrace,
    EditTestTrace,
    AddLogTrace,
    EditLogTrace,
    RemoveTraceArtifact,
    AddDocRef,
    EditDocRef,
    AddSrcCodeRef,
    EditSrcCodeRef,
    RemoveDataArtifact,
    LinkArtifactsComponents,
    RemoveArtComponentLink
)
from .change_result import (
    ProcessChangeResult,
    ProcessChangeError,
    ProcessChangeSuccess,
)

class State(BaseModel):
    """
    SpecLogician state instance container
    """
    file_location : str | None = Field(default=None)
    instances : list[StateInstance] = [StateInstance()]
    stash : list[StateInstance] = Field(default_factory=list[StateInstance]) # we'll use this to store discarded instances

    def curr_state(self) -> StateInstance:
        """  """
        return self.instances[0]

    def inst_list(self):
        """ Return a InstancesList (for pretty printing) """
        return InstancesList(states=self.instances)
    
    def process_change(
        self,
        change: StateChange,
        json_only: bool = False,
    ) -> ProcessChangeResult:
        """
        Apply a change and return a structured result.

        - json_only=False: returns object payload for rich printing / internal callers.
        - json_only=True : returns JSON-safe payload (summaries + json dumps).
        """

        old_state: StateInstance = self.curr_state()

        # Create a new candidate instance (deterministic snapshot + change list)
        state_copy = old_state.model_copy(deep=True)
        state_copy.update_created_at()
        state_copy.changes = [change]

        # 1) Apply change to state
        try:
            new_state: StateInstance = self.run_change_command(state_copy, change)
        except Exception as e:
            # recommendations still useful on failure (triage)
            failure = ChangeFailure(
                change=change.__class__.__name__,
                stage="apply_change",
                error_type=type(e).__name__,
                error=str(e),
            )
            # NOTE: no new_state exists; we recommend against old_state
            recs = Recommender().recommend(old_state, old_state, failure=failure)

            return ProcessChangeError(
                ok=False,
                stage="apply_change",
                error=e,
                change=change,
                old_state=old_state,
                # If your ProcessChangeError supports it, include recs:
                # recommendations=recs,
            )

        # 2) Perform logical analysis
        try:
            analyze_change(
                spec=new_state.spec,
                art_map=new_state.art_map,
                art_cont=new_state.art_container,
                change=change,
                json_only=json_only,
            )
        except Exception as e:
            failure = ChangeFailure(
                change=change.__class__.__name__,
                stage="analyze_change",
                error_type=type(e).__name__,
                error=str(e),
            )
            recs = Recommender().recommend(old_state, new_state, failure=failure)
            # Optionally store on the candidate for debugging:
            new_state.recommendations = recs

            return ProcessChangeError(
                ok=False,
                stage="analyze_change",
                error=e,
                change=change,
                old_state=old_state,
                # If supported:
                # new_state=new_state,
                # recommendations=recs,
            )

        # 3) Update stats (should not crash the whole pipeline ideally, but keeping your behavior)
        try:
            new_state.art_container.refresh_stats(new_state.art_map)
            refresh_spec_stats(new_state.spec, art_map=new_state.art_map)

        except Exception as e:
            failure = ChangeFailure(
                change=change.__class__.__name__,
                stage="refresh_stats",
                error_type=type(e).__name__,
                error=str(e),
            )
            recs = Recommender().recommend(old_state, new_state, failure=failure)
            new_state.recommendations = recs

            return ProcessChangeError(
                ok=False,
                stage="refresh_stats",
                error=e,
                change=change,
                old_state=old_state,
                # new_state=new_state,
            )

        # 3b) Trace reachability analysis (best-effort, derived artifact)
        try:
            if new_state.spec.scenario_reachability is not None:
                new_state.trace_reachability = compute_trace_reachability_analysis(
                    spec=new_state.spec,
                    art_cont=new_state.art_container,
                    art_map=new_state.art_map,
                    reach=new_state.spec.scenario_reachability,  # or wherever you store it
                    include_shortest_path=False,               # keep light by default
                )
        except Exception as e:
            return ProcessChangeError(
                ok=False,
                stage="trace_reachability",
                error=e,
                change=change,
                old_state=old_state,
            )

        # 4) Commit (newest at index 0)
        self.instances.insert(0, new_state)

        # 5) Diff (old -> new), only if we actually have a previous instance
        try:
            if len(self.instances) >= 2:
                # old = self.instances[1], new = self.instances[0]
                self.instances[0].state_diff = calc_diff(self.instances[1], self.instances[0])
            else:
                self.instances[0].state_diff = None
        except Exception as e:
            failure = ChangeFailure(
                change=change.__class__.__name__,
                stage="calc_diff",
                error_type=type(e).__name__,
                error=str(e),
            )
            # We *do* have a committed new state now
            recs = Recommender().recommend(old_state, self.instances[0], failure=failure)
            self.instances[0].recommendations = recs

            return ProcessChangeError(
                ok=False,
                stage="calc_diff",
                error=e,
                change=change,
                old_state=old_state,
                # new_state=self.instances[0],
            )

        # 6) Compute recommendations (success path)
        try:
            # old is previous instance if it exists; otherwise old_state (same as new baseline)
            prev = self.instances[1] if len(self.instances) >= 2 else self.instances[0]
            self.instances[0].recommendations = Recommender().recommend(prev, self.instances[0])
        except Exception as e:
            # Recommender should never take down the CLI; but if it does, degrade gracefully.
            self.instances[0].recommendations = [
                Recommendation(
                    text=f"Recommender failed: {type(e).__name__}: {e}",
                    kind="warning",
                    priority=5,
                )
            ]

        # 7) Success
        return ProcessChangeSuccess(
            ok=True,
            stage="ok",
            change=change,
            old_state=(self.instances[1] if len(self.instances) >= 2 else old_state),
            new_state=self.instances[0],
            state_diff=self.instances[0].state_diff,
            new_num_instances=len(self.instances),
            # If your success model includes it:
            # recommendations=self.instances[0].recommendations,
        )

    def run_change_command (
            self,
            s : StateInstance,
            change : StateChange,
        ) -> StateInstance:
        """ 
        Process change
        """

        if isinstance (change, DomainModelBaseEdit):
            s.spec.domain_model.set_base(change.new_base_src)

        elif isinstance (change, PredicateAdd):
            if s.spec.domain_model.pred_exists(change.pred_name):
                raise ValueError (f"Predicate with name {change.pred_name} already exists")            

            s.spec.domain_model.add_predicate(change.pred_name, change.pred_type, change.pred_src)

        elif isinstance (change, PredicateEdit):
            if not s.spec.domain_model.pred_exists(change.pred_name):
                raise ValueError (f"Predicate with name {change.pred_name} doesn't exist")

            s.spec.domain_model.edit_predicate(change.pred_name, change.pred_src)

        elif isinstance (change, PredicateRemove):
            if not s.spec.domain_model.pred_exists(change.pred_name):
                raise ValueError (f"Predicate with name {change.pred_name} doesn't exist")
            
            s.spec.domain_model.rem_predicate(change.pred_name)

        elif isinstance (change, TransitionAdd):
            if s.spec.domain_model.trans_exists(change.trans_name):
                raise ValueError (f"Transition with this name {change.trans_name} already exists!")

            s.spec.domain_model.add_transition(name=change.trans_name, src_code=change.trans_src)

        elif isinstance (change, TransitionEdit):
            if not s.spec.domain_model.trans_exists(change.trans_name):
                raise ValueError (f"Transition with this name {change.trans_name} doesn't exist")
            s.spec.domain_model.edit_transition(name=change.trans_name, src_code=change.trans_src)

        elif isinstance (change, TransitionRemove):
            if not s.spec.domain_model.trans_exists(change.trans_name):
                raise ValueError (f"Transition with this name {change.trans_name} doesn't exist")
            
            s.spec.domain_model.rem_transition(change.trans_name)

        elif isinstance(change, ScenarioAdd):
            if s.spec.scenario_exists(change.scenario_name):
                raise ValueError(f"Scenario with this name {change.scenario_name} already exists")

            s.spec.add_scenario(
                change.scenario_name,
                change.given,
                change.when,
                change.then,
                change.is_initial
            )

        elif isinstance (change, ScenarioEdit):
            if not s.spec.scenario_exists(change.scenario_name):
                raise ValueError (f"Scenario with the name '{change.scenario_name}' does not exist!")
            
            s.spec.edit_scenario(
                change.scenario_name, 
                change.given, 
                change.when, 
                change.then,
                change.is_initial
            )

        elif isinstance (change, ScenarioRemove):
            if not s.spec.scenario_exists(change.scenario_name):
                raise ValueError (f"Scenario with the name '{change.scenario_name}' does not exist!")

            s.spec.rem_scenario(change.scenario_name)

        # Test Traces

        elif isinstance(change, AddTestTrace):
            tr = TestTrace(
                art_id=change.art_id,
                # TraceArtifact core
                given=change.given,
                when=change.when,
                then=change.then,
                time=change.time,
                # TestTrace specific
                name=change.name,
                filepath=change.filepath,
                language=change.language,
                contents=change.contents,
            )
            s.art_container.add_trace(tr)

        elif isinstance(change, EditTestTrace):
            # Use model_fields_set so passing None explicitly can clear fields.
            patch: dict[str, object] = {}
            for f in change.model_fields_set:
                if f == "art_id":
                    continue
                patch[f] = getattr(change, f)
            s.art_container.edit_trace(change.art_id, patch)

        elif isinstance(change, AddLogTrace):
            tr = LogTrace(
                art_id=change.art_id,
                # TraceArtifact core
                given=change.given,
                when=change.when,
                then=change.then,
                time=change.time,
                # LogTrace specific
                filename=change.filename,
                contents=change.contents,
            )
            s.art_container.add_trace(tr)

        elif isinstance(change, EditLogTrace):
            patch: dict[str, object] = {}
            for f in change.model_fields_set:
                if f == "art_id":
                    continue
                patch[f] = getattr(change, f)
            s.art_container.edit_trace(change.art_id, patch)

        elif isinstance(change, RemoveTraceArtifact):
            s.art_container.rem_trace(change.art_id)

        # --- Data artifact changes (DocRef / SrcCodeRef) -----------------------------

        elif isinstance(change, AddDocRef):
            # Build DocRef and add via container API
            s.art_container.add_data_art(
                DocRef(
                    art_id=change.art_id,
                    text=change.text,
                    meta=change.meta,
                )
            )

        elif isinstance(change, EditDocRef):
            # Patch-edit via container API (only update provided fields)
            patch: dict[str, object] = {}
            if change.text is not None:
                patch["text"] = change.text
            if change.meta is not None:
                patch["meta"] = change.meta

            if patch:
                s.art_container.edit_data_art(change.art_id, patch)

        elif isinstance(change, AddSrcCodeRef):
            # Build SrcCodeRef and add via container API
            s.art_container.add_data_art(
                SrcCodeRef(
                    art_id=change.art_id,
                    src_code=change.src_code,
                    language=change.language,
                    file_path=change.file_path,
                    iml_code=change.iml_code,
                    meta=change.meta,
                )
            )

        elif isinstance(change, EditSrcCodeRef):
            # Patch-edit via container API (only update provided fields)
            patch: dict[str, object] = {}
            if change.src_code is not None:
                patch["src_code"] = change.src_code
            if change.language is not None:
                patch["language"] = change.language
            if change.file_path is not None:
                patch["file_path"] = change.file_path
            if change.iml_code is not None:
                patch["iml_code"] = change.iml_code
            if change.meta is not None:
                patch["meta"] = change.meta

            if patch:
                s.art_container.edit_data_art(change.art_id, patch)

        elif isinstance(change, RemoveDataArtifact):
            s.art_container.rem_data_art(change.art_id)

        elif isinstance (change, LinkArtifactsComponents):
            s.art_map.add_connection(art_id=change.art_id, comp_name=change.comp_name)

        elif isinstance (change, RemoveArtComponentLink):
            s.art_map.rem_connection(artifact_id=change.art_id, comp_name=change.comp_name)

        else:
            raise Exception(f"Unrecognized state change: {type(change).__name__}")

        return s

    def to_json(self):
        """ """
        return self.model_dump_json()

    @staticmethod
    def from_json(j : str):
        """
        Load in the State value
        """
        return State.model_validate_json(j)
    
    def save(self, dirpath : str | None | Path = None):
        """
        Save State to the specified directory
        """

        if dirpath is None:
            if self.file_location is None:
                dirpath = os.getcwd()
                filepath = os.path.join(dirpath, 'sl_state.json')
                console.print(f"Saving the state in CWD: {dirpath}")
            else:
                filepath = self.file_location
        else:
            filepath = os.path.join(dirpath, 'sl_state.json')
        
        with open(filepath, 'w') as outfile:
            print(self.to_json(), file=outfile)

    def save_to_path (self, path : Path) -> None:
        """ Save State to a specific filepath """
        try:
            with open(path, 'w') as outfile:
                print(self.to_json(), file=outfile)
        except Exception as e:
            raise ValueError(f"Failed to write state to disk [path = {path}]: {e}")

    @staticmethod
    def from_dir (dirpath : str | Path):
        """
        Load State from specified directory
        """
        try:
            state_path = os.path.join(dirpath, "sl_state.json")
            data = Path(state_path).read_text()
            state = State.from_json(data)
            state.file_location = state_path
        except Exception as e:
            console.print(f"Encountered error: {e}")
            return None
        
        return state
