#
#   Imandra Inc.
#
#   speclogician/utils/load.py
#

import json
import os
from pathlib import Path

import typer
from typing import NoReturn
from rich.prompt import Prompt

from ..state.state import State
from .__init__ import console


def _emit_json(obj: dict, exit_code: int = 0) -> NoReturn:
    """Print JSON and exit (useful for CLI / tests)."""
    typer.echo(json.dumps(obj, indent=2))
    raise typer.Exit(code=exit_code)


def load_state(path: Path | None = None, json: bool = False) -> State:
    """
    Load a State either from an explicit JSON path or from the local directory.

    - If `path` is provided:
        * if it exists: load JSON from it
        * if it doesn't exist: ask to create it (unless json)
    - If `path` is None:
        * load from current directory via State.from_dir()
        * if missing: ask to create it (unless json)

    If json=True:
      - emit JSON and Exit on errors (no prompts).
      - on success, return the state silently (no output).
    """
    cwd = Path(os.getcwd())

    def create_new_state(target: Path | None) -> State:
        """
        Create a new state and persist it.
        If target is a file path, save to that exact path.
        If target is None, save to cwd/sl_state.json.
        """
        state = State()

        if target is None:
            state.save(str(cwd))
            return state

        target = target.expanduser().resolve()
        target.parent.mkdir(parents=True, exist_ok=True)
        state.file_location = str(target)
        state.save_to_path(target)
        return state

    # -------------------------
    # Case 1: explicit file path
    # -------------------------
    if path is not None:
        path = path.expanduser()

        if path.exists():
            try:
                if path.is_dir():
                    new_state = State.from_dir(dirpath=path)
                else:
                    new_state = State.from_json(path.read_text())
                    new_state.file_location = str(path.resolve())
            except Exception as e:
                if json:
                    _emit_json(
                        {
                            "ok": False,
                            "error": "failed_to_load_state_json",
                            "path": str(path),
                            "message": str(e),
                        },
                        exit_code=2,
                    )
                console.print(f"[red]Failed to load state JSON:[/] {path}\n{e}")
                raise

            return new_state

        # File doesn't exist — auto-create in JSON mode (non-interactive)
        if json:
            return create_new_state(path)

        answer = Prompt.ask(
            prompt=f"⚠️ State file not found at {path}. Create a new one?",
            choices=["Yes", "No"],
            console=console,
            case_sensitive=False,
            default="Yes",
        )
        if answer.lower() == "no":
            console.print("Goodbye!")
            raise typer.Exit(0)

        console.print("Creating a new state!")
        return create_new_state(path)

    # -------------------------
    # Case 2: default local dir
    # -------------------------
    try:
        new_state = State.from_dir(dirpath=str(cwd))
    except Exception as e:
        if json:
            _emit_json(
                {
                    "ok": False,
                    "error": "failed_to_load_state_from_dir",
                    "dir": str(cwd),
                    "message": str(e),
                },
                exit_code=2,
            )
        console.print(f"[red]Failed to load state from directory:[/] {cwd}\n{e}")
        raise

    if new_state is not None:
        return new_state

    # Not found in cwd — auto-create in JSON mode (non-interactive)
    if json:
        return create_new_state(None)

    answer = Prompt.ask(
        prompt="⚠️ Couldn't find a state in current directory. Create a new one?",
        choices=["Yes", "No"],
        console=console,
        case_sensitive=False,
        default="Yes",
    )
    if answer.lower() == "no":
        console.print("Goodbye!")
        raise typer.Exit(0)

    console.print("Creating a new state!")
    return create_new_state(None)
