#
#   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, analyze_batch
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,
    ProcessBatchResult,
    ProcessBatchError,
    ProcessBatchSuccess,
)

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 undo(self, steps: int = 1) -> int:
        """
        Roll back `steps` state instances, moving them to the stash.

        Returns the number of instances actually undone (may be less than
        requested if there aren't enough instances to pop).

        Raises ValueError if there are no instances to undo (only one left).
        """
        if len(self.instances) <= 1:
            raise ValueError("Cannot undo: only one state instance exists (the initial state)")
        if steps < 1:
            raise ValueError("Steps must be >= 1")

        actually_undone = 0
        for _ in range(steps):
            if len(self.instances) <= 1:
                break
            removed = self.instances.pop(0)
            self.stash.append(removed)
            actually_undone += 1

        return actually_undone
    
    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:
                # new = self.instances[0], old = self.instances[1]
                self.instances[0].state_diff = calc_diff(self.instances[0], self.instances[1])
            else:
                self.instances[0].state_diff = None
        except Exception as e:
            # Roll back the commit so error and state are consistent
            self.instances.pop(0)

            failure = ChangeFailure(
                change=change.__class__.__name__,
                stage="calc_diff",
                error_type=type(e).__name__,
                error=str(e),
            )
            recs = Recommender().recommend(old_state, old_state, failure=failure)

            return ProcessChangeError(
                ok=False,
                stage="calc_diff",
                error=e,
                change=change,
                old_state=old_state,
            )

        # 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 process_batch(
        self,
        changes: list[StateChange],
        json_only: bool = False,
    ) -> ProcessBatchResult:
        """
        Apply multiple changes as a single batch.

        All mutations are applied first; analysis, stats, diff, and
        recommendations run once at the end.
        """
        if not changes:
            raise ValueError("Batch must contain at least one change")

        old_state: StateInstance = self.curr_state()

        # Create a single new candidate instance
        state_copy = old_state.model_copy(deep=True)
        state_copy.update_created_at()
        state_copy.changes = list(changes)

        # 1) Apply all changes sequentially
        for i, change in enumerate(changes):
            try:
                state_copy = self.run_change_command(state_copy, change)
            except Exception as e:
                failure = ChangeFailure(
                    change=change.__class__.__name__,
                    stage="apply_change",
                    error_type=type(e).__name__,
                    error=str(e),
                )
                recs = Recommender().recommend(old_state, old_state, failure=failure)
                return ProcessBatchError(
                    ok=False,
                    stage="apply_change",
                    error=e,
                    changes=changes,
                    change_index=i,
                    old_state=old_state,
                )

        new_state = state_copy

        # 2) Perform logical analysis (once for the whole batch)
        try:
            analyze_batch(
                spec=new_state.spec,
                art_cont=new_state.art_container,
                art_map=new_state.art_map,
                changes=changes,
                json_only=json_only,
            )
        except Exception as e:
            failure = ChangeFailure(
                change="batch",
                stage="analyze_change",
                error_type=type(e).__name__,
                error=str(e),
            )
            recs = Recommender().recommend(old_state, new_state, failure=failure)
            new_state.recommendations = recs
            return ProcessBatchError(
                ok=False,
                stage="analyze_change",
                error=e,
                changes=changes,
                change_index=-1,
                old_state=old_state,
            )

        # 3) Update stats
        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="batch",
                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 ProcessBatchError(
                ok=False,
                stage="refresh_stats",
                error=e,
                changes=changes,
                change_index=-1,
                old_state=old_state,
            )

        # 3b) Trace reachability analysis (best-effort)
        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,
                    include_shortest_path=False,
                )
        except Exception as e:
            return ProcessBatchError(
                ok=False,
                stage="trace_reachability",
                error=e,
                changes=changes,
                change_index=-1,
                old_state=old_state,
            )

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

        # 5) Diff (new = self.instances[0], old = self.instances[1])
        try:
            if len(self.instances) >= 2:
                self.instances[0].state_diff = calc_diff(self.instances[0], self.instances[1])
            else:
                self.instances[0].state_diff = None
        except Exception as e:
            # Roll back the commit so error and state are consistent
            self.instances.pop(0)

            failure = ChangeFailure(
                change="batch",
                stage="calc_diff",
                error_type=type(e).__name__,
                error=str(e),
            )
            recs = Recommender().recommend(old_state, old_state, failure=failure)
            return ProcessBatchError(
                ok=False,
                stage="calc_diff",
                error=e,
                changes=changes,
                change_index=-1,
                old_state=old_state,
            )

        # 6) Recommendations
        try:
            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:
            self.instances[0].recommendations = [
                Recommendation(
                    text=f"Recommender failed: {type(e).__name__}: {e}",
                    kind="warning",
                    priority=5,
                )
            ]

        # 7) Success
        return ProcessBatchSuccess(
            ok=True,
            stage="ok",
            changes=changes,
            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),
        )

    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
