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

from __future__ import annotations

import json
import typer
from enum import Enum
from datetime import datetime
from typing import NoReturn, Any
from rich.console import Console

from speclogician.state.change import StateChange
from speclogician.state.state import State
from speclogician.state.change_result import (
    ProcessChangeResult,
    ProcessChangeError,
    ProcessChangeSuccess,
    ProcessBatchResult,
    ProcessBatchError,
    ProcessBatchSuccess,
)

from speclogician.presentation.ctx import RenderCtx
from speclogician.presentation.builders.state_instance import build_state_instance_pm
from speclogician.presentation.builders.state_diff import build_state_diff_pm

from speclogician.presentation.renderers.state_instance import render_state_instance
from speclogician.presentation.renderers.state_diff import render_state_diff  # assuming you have it

console = Console()


def _dump_json(obj: Any) -> str:
    def _default(o: Any):
        if hasattr(o, "model_dump"):
            return o.model_dump(mode="json")
        if isinstance(o, Enum):
            return o.value
        if isinstance(o, datetime):
            return o.isoformat()
        raise TypeError(f"Object of type {type(o).__name__} is not JSON serializable")

    return json.dumps(
        obj,
        default=_default,
        indent=2,
        sort_keys=True,
    )

def run_change_and_render(
    st: State,
    change: StateChange,
    *,
    json_only: bool,
) -> NoReturn:
    """
    Unified runner for CLI commands.

    - json_only=True: emits JSON PM payload to stdout (stable, parseable).
    - json_only=False: pretty-prints PMs via Rich renderers.
    """

    # Always take the object path so we can drive presentation from PMs.
    result_obj: ProcessChangeResult = st.process_change(change, json_only=json_only)
    if result_obj["ok"]:
        st.save()

    # -------------------------
    # Error branch
    # -------------------------
    if not result_obj["ok"]:
        err_res: ProcessChangeError = result_obj

        # Build ctx from the old state if possible
        old_state = err_res["old_state"]
        ctx = RenderCtx(
            art_cont=old_state.art_container,
            art_map=old_state.art_map,
        )

        old_pm = build_state_instance_pm(old_state, ctx=ctx)

        if json_only:
            payload = {
                "ok": False,
                "stage": err_res["stage"],
                "change": err_res["change"].model_dump(mode="json"),
                "error": {
                    "type": type(err_res["error"]).__name__,
                    "message": str(err_res["error"]),
                },
                "old_state": old_pm.model_dump(mode="json"),
            }
            typer.echo(_dump_json(payload))
            raise typer.Exit(code=1)

        # human output
        console.print(f"[bold red]Error[/bold red] stage={err_res['stage']}")
        console.print(f"[red]{type(err_res['error']).__name__}: {err_res['error']}[/red]\n")
        console.print("[bold]Old state[/bold]")
        console.print(render_state_instance(old_pm, ctx=ctx))
        raise typer.Exit(code=1)

    # -------------------------
    # Success branch
    # -------------------------
    ok_res: ProcessChangeSuccess = result_obj

    old_state = ok_res["old_state"]
    new_state = ok_res["new_state"]
    sd = ok_res["state_diff"]

    # Prefer new-state context for presentation (has freshest art_map/art_cont)
    ctx = RenderCtx(
        art_cont=new_state.art_container,
        art_map=new_state.art_map,
    )

    old_pm = build_state_instance_pm(old_state, ctx=ctx)
    new_pm = build_state_instance_pm(new_state, ctx=ctx)
    diff_pm = build_state_diff_pm(sd, ctx=ctx)

    if json_only:
        payload = {
            "ok": True,
            "stage": "ok",
            "change": ok_res["change"].model_dump(mode="json"),
            "old_state": old_pm.model_dump(mode="json"),
            "new_state": new_pm.model_dump(mode="json"),
            "state_diff": diff_pm.model_dump(mode="json"),
            "new_num_instances": ok_res["new_num_instances"],
        }
        typer.echo(_dump_json(payload))
        raise typer.Exit(code=0)

    # human output
    console.print("[bold]Old state[/bold]")
    console.print(render_state_instance(old_pm, ctx=ctx))
    console.print()

    console.print("[bold]New state[/bold]")
    console.print(render_state_instance(new_pm, ctx=ctx))
    console.print()

    console.print("[bold]Diff[/bold]")
    console.print(render_state_diff(diff_pm, ctx=ctx))
    console.print()

    raise typer.Exit(code=0)


def run_batch_and_render(
    st: State,
    changes: list[StateChange],
    *,
    json_only: bool,
) -> NoReturn:
    """
    Unified runner for batch CLI commands.
    """

    result_obj: ProcessBatchResult = st.process_batch(changes, json_only=json_only)
    if result_obj["ok"]:
        st.save()

    # -------------------------
    # Error branch
    # -------------------------
    if not result_obj["ok"]:
        err_res: ProcessBatchError = result_obj

        old_state = err_res["old_state"]
        rctx = RenderCtx(
            art_cont=old_state.art_container,
            art_map=old_state.art_map,
        )

        old_pm = build_state_instance_pm(old_state, ctx=rctx)

        if json_only:
            payload = {
                "ok": False,
                "stage": err_res["stage"],
                "change_index": err_res["change_index"],
                "changes": [c.model_dump(mode="json") for c in err_res["changes"]],
                "error": {
                    "type": type(err_res["error"]).__name__,
                    "message": str(err_res["error"]),
                },
                "old_state": old_pm.model_dump(mode="json"),
            }
            typer.echo(_dump_json(payload))
            raise typer.Exit(code=1)

        # human output
        idx = err_res["change_index"]
        idx_str = f" (change #{idx})" if idx >= 0 else ""
        console.print(f"[bold red]Batch Error[/bold red] stage={err_res['stage']}{idx_str}")
        console.print(f"[red]{type(err_res['error']).__name__}: {err_res['error']}[/red]\n")
        console.print("[bold]Old state[/bold]")
        console.print(render_state_instance(old_pm, ctx=rctx))
        raise typer.Exit(code=1)

    # -------------------------
    # Success branch
    # -------------------------
    ok_res: ProcessBatchSuccess = result_obj

    old_state = ok_res["old_state"]
    new_state = ok_res["new_state"]
    sd = ok_res["state_diff"]

    rctx = RenderCtx(
        art_cont=new_state.art_container,
        art_map=new_state.art_map,
    )

    old_pm = build_state_instance_pm(old_state, ctx=rctx)
    new_pm = build_state_instance_pm(new_state, ctx=rctx)
    diff_pm = build_state_diff_pm(sd, ctx=rctx)

    if json_only:
        payload = {
            "ok": True,
            "stage": "ok",
            "changes": [c.model_dump(mode="json") for c in ok_res["changes"]],
            "old_state": old_pm.model_dump(mode="json"),
            "new_state": new_pm.model_dump(mode="json"),
            "state_diff": diff_pm.model_dump(mode="json"),
            "new_num_instances": ok_res["new_num_instances"],
        }
        typer.echo(_dump_json(payload))
        raise typer.Exit(code=0)

    # human output
    console.print(f"[bold]Batch applied ({len(ok_res['changes'])} changes)[/bold]\n")

    console.print("[bold]Old state[/bold]")
    console.print(render_state_instance(old_pm, ctx=rctx))
    console.print()

    console.print("[bold]New state[/bold]")
    console.print(render_state_instance(new_pm, ctx=rctx))
    console.print()

    console.print("[bold]Diff[/bold]")
    console.print(render_state_diff(diff_pm, ctx=rctx))
    console.print()

    raise typer.Exit(code=0)