#
#   Imandra Inc.
#
#   speclogician/shell/ch_shell.py
#

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Optional

from cmd2 import Cmd, with_argparser

from speclogician.modeling.predicates import PredicateType
from speclogician.modeling.scenario import ScenarioDelta
from speclogician.state.change import (
    DomainModelBaseEdit,
    PredicateAdd,
    PredicateEdit,
    PredicateRemove,
    TransitionAdd,
    TransitionEdit,
    TransitionRemove,
    ScenarioAdd,
    ScenarioEdit,
    ScenarioRemove,
    # later: AddTestTrace, EditTestTrace, AddLogTrace, ...
)
from speclogician.utils.load import load_state


def _print_result(cmd: Cmd, res: object) -> None:
    """
    Best-effort printer for ProcessChangeResult / ProcessChangeObjectResult.
    - dict: dump as JSON
    - pydantic-ish: model_dump
    - fallback: str()
    """
    if isinstance(res, dict):
        cmd.poutput(json.dumps(res, indent=2, default=str))
        return
    if hasattr(res, "model_dump"):
        cmd.poutput(json.dumps(res.model_dump(mode="json"), indent=2, default=str))  # type: ignore[attr-defined]
        return
    cmd.poutput(str(res))


def _split_csv(s: Optional[str]) -> list[str]:
    if not s:
        return []
    return [x.strip() for x in s.split(",") if x.strip()]


def _merge_dedup_preserve_order(*lists: list[str]) -> list[str]:
    out: list[str] = []
    seen: set[str] = set()
    for xs in lists:
        for x in xs:
            if x not in seen:
                out.append(x)
                seen.add(x)
    return out


def _delta_from_add(xs: list[str]) -> ScenarioDelta | None:
    d = ScenarioDelta(add=xs)
    return None if d.is_empty() else d


class ChangeShell(Cmd):
    """
    Nested cmd2 shell that exposes the same 'ch' operations as the CLI.
    """

    prompt = "sl ch> "
    intro = "Change commands (type 'help' to list, 'exit' to return)\n"
    allow_cli_args = False

    def __init__(self, parent: Cmd):
        super().__init__(allow_cli_args=False)
        self.parent = parent

    # -------------------------
    # shared accessors
    # -------------------------

    def _get_state(self):
        """
        Use parent.state if present, otherwise load state from default location.
        """
        st = getattr(self.parent, "state", None)
        if st is None:
            st = load_state()
            setattr(self.parent, "state", st)
        return st

    def _json(self) -> bool:
        return bool(getattr(self.parent, "json_default", False))

    # -------------------------
    # base-edit  (matches Typer: ch base-edit --src/--file)
    # -------------------------

    base_edit_parser = argparse.ArgumentParser()
    g = base_edit_parser.add_mutually_exclusive_group(required=True)
    g.add_argument("--src", help="IML source code for the domain base component")
    g.add_argument("--file", type=Path, help="Path to a file containing IML base source code")
    @with_argparser(base_edit_parser)
    def do_base_edit(self, args) -> None:
        st = self._get_state()

        src: str
        if args.src is not None:
            src = args.src
        else:
            p: Path = args.file
            src = p.read_text()

        change = DomainModelBaseEdit(new_base_src=src)
        res = st.process_change(change, json_only=self._json())
        _print_result(self, res)

    # -------------------------
    # predicates
    # -------------------------

    pred_add_parser = argparse.ArgumentParser()
    pred_add_parser.add_argument("--pred-type", choices=["state", "action"], required=True)
    pred_add_parser.add_argument("--pred-name", required=True)
    pred_add_parser.add_argument("--pred-src", required=True)
    @with_argparser(pred_add_parser)
    def do_pred_add(self, args) -> None:
        st = self._get_state()
        ptype = PredicateType.STATE if args.pred_type == "state" else PredicateType.ACTION
        change = PredicateAdd(pred_type=ptype, pred_name=args.pred_name, pred_src=args.pred_src)
        res = st.process_change(change, json_only=self._json())
        _print_result(self, res)

    pred_edit_parser = argparse.ArgumentParser()
    pred_edit_parser.add_argument("--pred-name", required=True)
    pred_edit_parser.add_argument("--pred-src", required=True)
    @with_argparser(pred_edit_parser)
    def do_pred_edit(self, args) -> None:
        st = self._get_state()
        change = PredicateEdit(pred_name=args.pred_name, pred_src=args.pred_src)
        res = st.process_change(change, json_only=self._json())
        _print_result(self, res)

    pred_remove_parser = argparse.ArgumentParser()
    pred_remove_parser.add_argument("--pred-name", required=True)
    @with_argparser(pred_remove_parser)
    def do_pred_remove(self, args) -> None:
        st = self._get_state()
        change = PredicateRemove(pred_name=args.pred_name)
        res = st.process_change(change, json_only=self._json())
        _print_result(self, res)

    # -------------------------
    # transitions
    # -------------------------

    trans_add_parser = argparse.ArgumentParser()
    trans_add_parser.add_argument("--trans-name", required=True)
    trans_add_parser.add_argument("--trans-src", required=True)
    @with_argparser(trans_add_parser)
    def do_trans_add(self, args) -> None:
        st = self._get_state()
        change = TransitionAdd(trans_name=args.trans_name, trans_src=args.trans_src)
        res = st.process_change(change, json_only=self._json())
        _print_result(self, res)

    trans_edit_parser = argparse.ArgumentParser()
    trans_edit_parser.add_argument("--trans-name", required=True)
    trans_edit_parser.add_argument("--trans-src", required=True)
    @with_argparser(trans_edit_parser)
    def do_trans_edit(self, args) -> None:
        st = self._get_state()
        change = TransitionEdit(trans_name=args.trans_name, trans_src=args.trans_src)
        res = st.process_change(change, json_only=self._json())
        _print_result(self, res)

    trans_remove_parser = argparse.ArgumentParser()
    trans_remove_parser.add_argument("--trans-name", required=True)
    @with_argparser(trans_remove_parser)
    def do_trans_remove(self, args) -> None:
        st = self._get_state()
        change = TransitionRemove(trans_name=args.trans_name)
        res = st.process_change(change, json_only=self._json())
        _print_result(self, res)

    # -------------------------
    # scenarios
    # -------------------------

    scenario_add_parser = argparse.ArgumentParser()
    scenario_add_parser.add_argument("scenario_name")
    scenario_add_parser.add_argument("--given", action="append", default=[], help="Repeatable")
    scenario_add_parser.add_argument("--when", action="append", default=[], help="Repeatable")
    scenario_add_parser.add_argument("--then", action="append", default=[], help="Repeatable")
    scenario_add_parser.add_argument("--given-list", default=None)
    scenario_add_parser.add_argument("--when-list", default=None)
    scenario_add_parser.add_argument("--then-list", default=None)
    @with_argparser(scenario_add_parser)
    def do_scenario_add(self, args) -> None:
        st = self._get_state()

        given_all = _merge_dedup_preserve_order(args.given, _split_csv(args.given_list))
        when_all = _merge_dedup_preserve_order(args.when, _split_csv(args.when_list))
        then_all = _merge_dedup_preserve_order(args.then, _split_csv(args.then_list))

        change = ScenarioAdd(
            scenario_name=args.scenario_name,
            given=_delta_from_add(given_all),
            when=_delta_from_add(when_all),
            then=_delta_from_add(then_all),
        )
        res = st.process_change(change, json_only=self._json())
        _print_result(self, res)

    scenario_edit_parser = argparse.ArgumentParser()
    scenario_edit_parser.add_argument("scenario_name")
    scenario_edit_parser.add_argument("--given", default=None)
    scenario_edit_parser.add_argument("--when", default=None)
    scenario_edit_parser.add_argument("--then", default=None)
    @with_argparser(scenario_edit_parser)
    def do_scenario_edit(self, args) -> None:
        st = self._get_state()

        # NOTE: to match CLI behavior, reuse the SAME parser you use in Typer.
        # If you already have parse_scenario_delta(), import it and use it here.
        from speclogician.commands.utils import parse_scenario_delta  # adjust import to your actual module

        change = ScenarioEdit(
            scenario_name=args.scenario_name,
            given=parse_scenario_delta(args.given),
            when=parse_scenario_delta(args.when),
            then=parse_scenario_delta(args.then),
        )
        res = st.process_change(change, json_only=self._json())
        _print_result(self, res)

    scenario_remove_parser = argparse.ArgumentParser()
    scenario_remove_parser.add_argument("scenario_name")
    @with_argparser(scenario_remove_parser)
    def do_scenario_remove(self, args) -> None:
        st = self._get_state()
        change = ScenarioRemove(scenario_name=args.scenario_name)
        res = st.process_change(change, json_only=self._json())
        _print_result(self, res)

    # -------------------------
    # navigation
    # -------------------------

    def do_exit(self, _arg) -> bool:
        """Return to the main shell."""
        return True

    def do_EOF(self, _arg) -> bool:
        """Ctrl-D exits this sub-shell."""
        return True