#
#   Imandra Inc.
#
#   speclogician/commands/ch_cmd.py
#

from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import NoReturn, Optional

import typer

from ..state.change import (
    DomainModelBaseEdit,
    PredicateAdd,
    PredicateEdit,
    PredicateRemove,
    TransitionAdd,
    TransitionEdit,
    TransitionRemove,
    ScenarioAdd,
    ScenarioEdit,
    ScenarioRemove,
    parse_change,
)

from ..modeling.scenario import ScenarioDelta, BoolDelta
from ..modeling.predicates import PredicateType
from .utils import parse_scenario_delta, get_state
from .runner import run_change_and_render, run_batch_and_render
from .cmd_ch_trace import app as trace_app
from .cmd_ch_data import app as data_app


def _emit_parse_error(*, stage: str, message: str, change_index: int) -> NoReturn:
    """Emit a structured JSON error for pre-runner failures and exit."""
    payload = {
        "ok": False,
        "stage": stage,
        "change_index": change_index,
        "error": {
            "type": "ParseError",
            "message": message,
        },
    }
    typer.echo(json.dumps(payload, indent=2, sort_keys=True))
    raise typer.Exit(code=1)

app = typer.Typer(
    help=(
        "Apply exactly one change to the current SpecLogician state.\n\n"
        "Each change produces a new immutable state, triggers analysis, and writes a state diff.\n"
        "Use `--json` for machine-readable output."
    )
)

app.add_typer(trace_app, name="trace")
app.add_typer(data_app, name="data")

# -----------------------------------------------------------------------------
# Batch
# -----------------------------------------------------------------------------

@app.command(
    "batch",
    help=(
        "Apply multiple changes in a single batch.\n\n"
        "Changes are applied sequentially, but analysis runs only once at the end.\n"
        "Input is a JSON array of change objects, each with a 'kind' discriminator.\n\n"
        "Reads from --file or stdin."
    ),
)
def batch_cmd(
    ctx: typer.Context,
    file: Optional[Path] = typer.Option(
        None,
        "--file",
        exists=True,
        readable=True,
        help="Path to a JSON file containing an array of change objects.",
    ),
    json_only: bool = typer.Option(
        False,
        "--json",
        help="Emit JSON only (no rich output).",
    ),
    dry_run: bool = typer.Option(
        False,
        "--dry-run",
        help="Parse and validate all changes without applying them or running analysis.",
    ),
) -> None:
    """
    Examples:

      speclogician ch batch --file changes.json

      speclogician ch batch --json < changes.json

      speclogician ch batch --dry-run --json < changes.json

      echo '[{"kind":"PredicateAdd","pred_type":"state","pred_name":"x_gt_0","pred_src":"s.x > 0"}]' | speclogician ch batch
    """
    state = get_state(ctx, json_mode=json_only)

    # Read input
    if file is not None:
        raw = file.read_text()
    else:
        raw = sys.stdin.read()

    try:
        data = json.loads(raw)
    except json.JSONDecodeError as e:
        if json_only:
            _emit_parse_error(stage="json_parse", message=f"Invalid JSON: {e}", change_index=-1)
        raise typer.BadParameter(f"Invalid JSON: {e}")

    if not isinstance(data, list):
        if json_only:
            _emit_parse_error(stage="input_validation", message="Expected a JSON array of change objects", change_index=-1)
        raise typer.BadParameter("Expected a JSON array of change objects")
    if len(data) == 0:
        if json_only:
            _emit_parse_error(stage="input_validation", message="Batch must contain at least one change", change_index=-1)
        raise typer.BadParameter("Batch must contain at least one change")

    changes = []
    for i, obj in enumerate(data):
        if not isinstance(obj, dict):
            msg = f"Change at index {i} is not a JSON object"
            if json_only:
                _emit_parse_error(stage="change_parse", message=msg, change_index=i)
            raise typer.BadParameter(msg)
        try:
            changes.append(parse_change(obj))
        except (ValueError, Exception) as e:
            msg = f"Change at index {i}: {e}"
            if json_only:
                _emit_parse_error(stage="change_parse", message=str(e), change_index=i)
            raise typer.BadParameter(msg)

    if dry_run:
        if json_only:
            payload = {
                "ok": True,
                "dry_run": True,
                "changes_parsed": len(changes),
                "changes": [c.model_dump(mode="json") for c in changes],
            }
            typer.echo(json.dumps(payload, indent=2, sort_keys=True))
        else:
            from ..utils import console
            console.print(f"[bold green]Dry run OK:[/bold green] {len(changes)} change(s) parsed successfully.")
            for i, c in enumerate(changes):
                console.print(f"  [{i}] {c.kind}")
        raise typer.Exit(code=0)

    run_batch_and_render(state, changes, json_only=json_only)
    state.save()


# -----------------------------------------------------------------------------
# Base Model
# -----------------------------------------------------------------------------

@app.command(
    "base-edit",
    help=(
        "Edit the domain base (types + core invariants) using IML source.\n\n"
        "Exactly one of --src or --file must be provided."
    ),
)
def domain_base_edit(
    ctx: typer.Context,
    src: Optional[str] = typer.Option(
        None,
        "--src",
        help="IML source code for the domain base component (types + invariants).",
    ),
    file: Optional[Path] = typer.Option(
        None,
        "--file",
        exists=True,
        readable=True,
        help="Path to a file containing IML base source code.",
    ),
    json_only: bool = typer.Option(
        False,
        "--json",
        help="Emit JSON only (no rich output).",
    ),
) -> None:
    """
    Examples:

      speclogician ch base-edit --src "type state = { x:int }\ntype action = int"

      speclogician ch base-edit --file ./domain_base.iml

      speclogician ch base-edit --file ./domain_base.iml --json
    """
    state = get_state(ctx, json_mode=json_only)

    if src is None and file is None:
        raise typer.BadParameter("One of --src or --file must be provided")
    if src is not None and file is not None:
        raise typer.BadParameter("Only one of --src or --file may be provided")
    if file is not None:
        src = file.read_text()

    assert src is not None
    change = DomainModelBaseEdit(new_base_src=src)
    run_change_and_render(state, change, json_only=json_only)
    state.save()


# -----------------------------------------------------------------------------
# Scenarios
# -----------------------------------------------------------------------------

@app.command(
    name="scenario-add",
    help=(
        "Add a new scenario (set-delta semantics).\n\n"
        "Use --given/--when/--then (repeatable) or the corresponding CSV flags.\n"
        "Use --initial/--no-initial to control whether this scenario is an initial root for reachability."
    ),
)
def scenario_add(
    ctx: typer.Context,
    scenario_name: str = typer.Argument(..., help="Name of the scenario to add."),
    given: list[str] = typer.Option(
        [],
        "--given",
        help=(
            "State predicate name (repeatable).\n"
            "Examples:\n"
            "  --given pred_eq_ten\n"
            "  --given pred_gte_ten --given pred_ne_ten"
        ),
    ),
    when: list[str] = typer.Option(
        [],
        "--when",
        help=(
            "Action predicate name (repeatable).\n"
            "Examples:\n"
            "  --when act_gt_x\n"
            "  --when act_gt_x --when act_lte_x"
        ),
    ),
    then: list[str] = typer.Option(
        [],
        "--then",
        help=(
            "Transition name (repeatable).\n"
            "Examples:\n"
            "  --then step1\n"
            "  --then step1 --then step2"
        ),
    ),
    given_list: str | None = typer.Option(
        None,
        "--given-list",
        help='Comma-separated list of state predicates. Example: --given-list "p1,p2,p3"',
    ),
    when_list: str | None = typer.Option(
        None,
        "--when-list",
        help='Comma-separated list of action predicates. Example: --when-list "q1,q2"',
    ),
    then_list: str | None = typer.Option(
        None,
        "--then-list",
        help='Comma-separated list of transitions. Example: --then-list "t1,t2"',
    ),
    initial: bool | None = typer.Option(
        None,
        "--initial/--no-initial",
        help="If set, marks this scenario as initial (reachability root).",
    ),
    json_only: bool = typer.Option(
        False,
        "--json",
        help="Emit JSON only (includes parsed change + state summaries + diff).",
    ),
) -> None:
    """
    Examples:

      speclogician ch scenario-add sc1 --given p1 --when q1 --then t1

      speclogician ch scenario-add boot --given init --then init_step --initial

      speclogician ch scenario-add sc2 --given-list "p1,p2" --then-list "t1,t2" --json
    """
    state = get_state(ctx, json_mode=json_only)

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

    def _merge(a: list[str], b: list[str]) -> list[str]:
        # dedup preserving order
        return list(dict.fromkeys([*a, *b]))

    given_all = _merge(given, _split_csv(given_list))
    when_all = _merge(when, _split_csv(when_list))
    then_all = _merge(then, _split_csv(then_list))

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

    is_initial = BoolDelta(set=initial) if initial is not None else None

    change = ScenarioAdd(
        scenario_name=scenario_name,
        given=_delta_from_add(given_all),
        when=_delta_from_add(when_all),
        then=_delta_from_add(then_all),
        is_initial=is_initial,
    )

    run_change_and_render(state, change, json_only=json_only)
    state.save()


@app.command(
    name="scenario-edit",
    help=(
        "Edit an existing scenario by adding/removing predicates/transitions.\n\n"
        "Deltas accept:\n"
        "  - '+a,+b,-c' (add/remove)\n"
        "  - 'a,b,c'    (add-only)\n"
        "  - JSON string or path to JSON file\n"
        "Use --initial/--no-initial to update whether the scenario is initial."
    ),
)
def scenario_edit_cmd(
    ctx: typer.Context,
    scenario_name: str = typer.Argument(..., help="Name of the scenario to edit."),
    given: str | None = typer.Option(
        None,
        "--given",
        help=(
            "Delta for GIVEN (state predicates). Example:\n"
            "  --given '+p1,-p2'\n"
            "  --given 'p1,p2'\n"
            "  --given @/path/to/delta.json"
        ),
    ),
    when: str | None = typer.Option(
        None,
        "--when",
        help=(
            "Delta for WHEN (action predicates). Example:\n"
            "  --when '+q1,-q2'\n"
            "  --when 'q1,q2'\n"
            "  --when @/path/to/delta.json"
        ),
    ),
    then: str | None = typer.Option(
        None,
        "--then",
        help=(
            "Delta for THEN (transition names). Example:\n"
            "  --then '+t1,-t2'\n"
            "  --then 't1,t2'\n"
            "  --then @/path/to/delta.json"
        ),
    ),
    initial: bool | None = typer.Option(
        None,
        "--initial/--no-initial",
        help="If set, update whether this scenario is initial (reachability root).",
    ),
    json_only: bool = typer.Option(
        False,
        "--json",
        help="Emit JSON only (includes parsed change + state summaries + diff).",
    ),
) -> None:
    """
    Examples:

      speclogician ch scenario-edit sc1 --given '+p1,-p2'

      speclogician ch scenario-edit boot --initial

      speclogician ch scenario-edit boot --no-initial --json
    """
    state = get_state(ctx, json_mode=json_only)
    is_initial = BoolDelta(set=initial) if initial is not None else None

    change = ScenarioEdit(
        scenario_name=scenario_name,
        given=parse_scenario_delta(given),
        when=parse_scenario_delta(when),
        then=parse_scenario_delta(then),
        is_initial=is_initial,
    )
    run_change_and_render(state, change, json_only=json_only)
    state.save()


@app.command(
    name="scenario-remove",
    help="Remove a scenario from the current spec/state by name.",
)
def scenario_remove_cmd(
    ctx: typer.Context,
    scenario_name: str = typer.Argument(
        ...,
        help="Name of the scenario to remove.",
    ),
    json_only: bool = typer.Option(
        False,
        "--json",
        help="Emit JSON only (includes change + state summaries + diff).",
    ),
) -> None:
    """
    Examples:

      speclogician ch scenario-remove sc1

      speclogician ch scenario-remove sc1 --json
    """
    state = get_state(ctx, json_mode=json_only)
    change = ScenarioRemove(scenario_name=scenario_name)
    run_change_and_render(state, change, json_only=json_only)
    state.save()




# -----------------------------------------------------------------------------
# Predicates
# -----------------------------------------------------------------------------

@app.command("pred-add", help="Add a new predicate to the domain model.")
def predicate_add(
    ctx: typer.Context,
    pred_type: str = typer.Option(
        ...,
        "--pred-type",
        help="Predicate type: 'state' or 'action'.",
    ),
    pred_name: str = typer.Option(..., "--pred-name", help="Predicate name (identifier)."),
    pred_src: str = typer.Option(..., "--pred-src", help="Predicate IML source code."),
    json_only: bool = typer.Option(False, "--json", help="Emit JSON only."),
) -> None:
    """
    Example:

      speclogician ch pred-add --pred-type state --pred-name x_eq_1 --pred-src "s.x = 1"
    """
    state = get_state(ctx, json_mode=json_only)

    pt = (pred_type or "").strip().lower()
    if pt not in ("state", "action"):
        raise typer.BadParameter("--pred-type must be one of: state, action")

    change = PredicateAdd(
        pred_type=PredicateType(pt),
        pred_name=pred_name,
        pred_src=pred_src,
    )
    run_change_and_render(state, change, json_only=json_only)
    state.save()


@app.command("pred-edit", help="Edit an existing predicate in the domain model.")
def predicate_edit(
    ctx: typer.Context,
    pred_name: str = typer.Option(..., "--pred-name", help="Predicate name to edit."),
    pred_src: str = typer.Option(..., "--pred-src", help="Updated predicate IML source."),
    json_only: bool = typer.Option(False, "--json", help="Emit JSON only."),
) -> None:
    state = get_state(ctx, json_mode=json_only)
    change = PredicateEdit(pred_name=pred_name, pred_src=pred_src)
    run_change_and_render(state, change, json_only=json_only)
    state.save()


@app.command("pred-remove", help="Remove a predicate from the domain model.")
def predicate_remove(
    ctx: typer.Context,
    pred_name: str = typer.Option(..., "--pred-name", help="Predicate name to remove."),
    json_only: bool = typer.Option(False, "--json", help="Emit JSON only."),
) -> None:
    state = get_state(ctx, json_mode=json_only)
    change = PredicateRemove(pred_name=pred_name)
    run_change_and_render(state, change, json_only=json_only)
    state.save()


# -----------------------------------------------------------------------------
# Transitions
# -----------------------------------------------------------------------------

@app.command("trans-add", help="Add a new transition function to the domain model.")
def transition_add(
    ctx: typer.Context,
    trans_name: str = typer.Option(..., "--trans-name", help="Transition name (identifier)."),
    trans_src: str = typer.Option(..., "--trans-src", help="Transition IML source code."),
    json_only: bool = typer.Option(False, "--json", help="Emit JSON only."),
) -> None:
    state = get_state(ctx, json_mode=json_only)
    change = TransitionAdd(trans_name=trans_name, trans_src=trans_src)
    run_change_and_render(state, change, json_only=json_only)
    state.save()


@app.command("trans-edit", help="Edit an existing transition function in the domain model.")
def transition_edit(
    ctx: typer.Context,
    trans_name: str = typer.Option(..., "--trans-name", help="Transition name to edit."),
    trans_src: str = typer.Option(..., "--trans-src", help="Updated transition IML source code."),
    json_only: bool = typer.Option(False, "--json", help="Emit JSON only."),
) -> None:
    state = get_state(ctx, json_mode=json_only)
    change = TransitionEdit(trans_name=trans_name, trans_src=trans_src)
    run_change_and_render(state, change, json_only=json_only)
    state.save()


@app.command("trans-remove", help="Remove a transition function from the domain model.")
def transition_remove(
    ctx: typer.Context,
    trans_name: str = typer.Option(..., "--trans-name", help="Transition name to remove."),
    json_only: bool = typer.Option(False, "--json", help="Emit JSON only."),
) -> None:
    state = get_state(ctx, json_mode=json_only)
    change = TransitionRemove(trans_name=trans_name)
    run_change_and_render(state, change, json_only=json_only)
    state.save()


# -----------------------------------------------------------------------------
# Artifact Mapping
# -----------------------------------------------------------------------------
