#
#   Imandra Inc.
#
#   speclogician/logic/strings.py
#


from __future__ import annotations

import re
from typing import Sequence

from ..modeling.scenario import Scenario

def parens(x: str) -> str:
    x = x.strip()
    return x if (x.startswith("(") and x.endswith(")")) else f"({x})"


def and_all(conjuncts: Sequence[str]) -> str:
    if not conjuncts:
        return "true"
    return " && ".join(parens(c) for c in conjuncts)


def mk_predicates_only_instance_src(
    name: str,
    *,
    given: Sequence[str],     # state predicate names: p : state -> bool
    when: Sequence[str],      # action predicate names: q : state -> action -> bool (state first!)
    state_type: str = "state",
    action_type: str = "action",
    state_var: str = "s",
    action_var: str = "a",
) -> str:
    """
    Consistency query: only predicates (given/when), no transitions.

    Produces:
          (p1 s) && (q1 s a) && ...
    """
    lines : list[str] = ["fun (s : state) (a : action) -> "]

    exprs: list[str] = []
    exprs += [f"{p} {state_var}" for p in given]
    exprs += [f"{q} {state_var} {action_var}" for q in when]

    lines.append(f"    {and_all(exprs)}")
    result =  "\n".join(lines)

    return result


def mk_same_action_trace_instance_src(
    name: str,
    *,
    given: Sequence[str],        # state predicate names
    when: Sequence[str],         # action predicate names
    transitions: Sequence[str],  # transition fn names: t : state -> action -> state
    state_type: str = "state",
    action_type: str = "action",
    state0_var: str = "s0",
    action_var: str = "a",
    when_each_step: bool = True,  # apply action predicates at each pre-state
) -> str:
    """
    Trace query: transitions applied sequentially using the SAME action `a`.

    Produces:

      instance <name> :
        let s0 : state = s0 in
        let a  : action = a  in
        let s1 = t0 s0 a in
        let s2 = t1 s1 a in
          (given(s0)) && (when(s0,a)) && (when(s1,a)) ...

    Note: for N transitions, `when` is applied to s0..s{N-1} if when_each_step=True.
    """
    lines: list[str] = [f"instance {name} :"]

    # typed existentials
    lines.append(f"  let {state0_var} : {state_type} = {state0_var} in")
    lines.append(f"  let {action_var} : {action_type} = {action_var} in")

    # let-chain states
    # s0 already exists; create s1..sN
    for i, tname in enumerate(transitions):
        si = state0_var if i == 0 else f"s{i}"
        si1 = f"s{i+1}"
        lines.append(f"  let {si1} = {tname} {si} {action_var} in")

    exprs: list[str] = []
    exprs += [f"{p} {state0_var}" for p in given]

    if when_each_step:
        for i in range(len(transitions)):
            si = state0_var if i == 0 else f"s{i}"
            exprs += [f"{q} {si} {action_var}" for q in when]
    else:
        exprs += [f"{q} {state0_var} {action_var}" for q in when]

    lines.append(f"    {and_all(exprs)}")
    return "\n".join(lines)


def mk_scenario_expr(sc: Scenario, *, s_var: str = "s", a_var: str = "a") -> str:
    parts: list[str] = []
    parts += [f"({p} {s_var})" for p in sc.given]
    parts += [f"({q} {s_var} {a_var})" for q in sc.when]
    return "true" if not parts else " && ".join(parts)

def mk_spec_disjunction_expr(scenarios: Sequence[Scenario]) -> str:
    disj = " || ".join(f"({mk_scenario_expr(sc)})" for sc in scenarios)
    return "false" if not disj else disj

def mk_goal_def(goal_name: str, goal_expr: str, *, needs_action: bool) -> str:
    if needs_action:
        return (
            f"let {goal_name} (s : state) (a : action) : bool =\n"
            f"  {goal_expr}\n"
        )
    else:
        return (
            f"let {goal_name} (s : state) : bool =\n"
            f"  {goal_expr}\n"
        )

def collect_pred_names(scenarios: Sequence[Scenario]) -> tuple[list[str], list[str]]:
    state_names: list[str] = []
    action_names: list[str] = []
    for sc in scenarios:
        state_names += list(sc.given)
        action_names += list(sc.when)
    return state_names, action_names

LET_RE = re.compile(r"^\s*let\s+([a-zA-Z_][a-zA-Z0-9_']*)\s*(?::\s*[^=]+)?=\s*(.+)$", re.DOTALL)

def normalize_optional(s: str | None) -> str | None:
    if s is None:
        return None
    s2 = s.strip()
    return s2 if s2 else None

def mk_typed_let(binding_src: str, *, expected_type: str, default_name: str) -> tuple[str, str]:
    """
    Returns: (let_src, bound_name)
    Accepts either:
      - `let x = ...` (or `let x : t = ...`)  -> keeps x
      - `<expr>`                              -> wraps as `let <default_name> : expected_type = (<expr>)`
    """
    m = LET_RE.match(binding_src.strip())
    if m:
        name = m.group(1)
        # Keep user's let as-is; we just trust their annotation (or lack of).
        # (Type errors will be caught by eval_model.)
        return binding_src.strip(), name

    name = default_name
    let_src = f"let {name} : {expected_type} = (\n{binding_src.strip()}\n)"
    return let_src, name

def mk_instance_src_for_expr(*, name: str, expr: str, needs_action: bool) -> str:
    # Assumes `state` and (optionally) `action` exist in the model context.
    # `expr` must be boolean and can mention `s` and optionally `a`.
    if needs_action:
        return f"""
instance {name} =
  {{ s : state; a : action }}
  where {expr}
"""
    return f"""
instance {name} =
  {{ s : state }}
  where {expr}
"""

def errs_to_string(eval_res: object) -> str | None:
    """
    Your eval_model result looks like it may have .has_errors and .errors.
    Convert to a multi-line string.
    """
    if not bool(getattr(eval_res, "has_errors", False)):
        return None

    errs = getattr(eval_res, "errors", None)
    if not errs:
        return "Domain/trace model has errors"

    # errs items could be strings or structured; stringify safely
    lines: list[str] = []
    for e in errs:
        lines.append(str(e))
    return "\n".join(lines)

def mk_conj(xs: list[str]) -> str:
    if not xs:
        return "true"
    return " && \n".join(f"({x})" for x in xs)

def mk_apply_then_expr(
        trans_names: list[str], 
        *, 
        s_var: str, 
        a_var: str
) -> str:
    """
    Build an IML expression for sequential transition application.
    Uses the same semantics as check_trace_match.
    """
    if not trans_names:
        return s_var
    if len(trans_names) == 1:
        return f"({trans_names[0]} {s_var} {a_var})"

    lets: list[str] = []
    prev = s_var
    for i, t in enumerate(trans_names, 1):
        si = f"{s_var}{i}"
        lets.append(f"let {si} = ({t} {prev} {a_var}) in")
        prev = si
    return " ".join(lets) + f" {prev}"

def mk_apply_then_let_block(
    then_names: list[str],
    *,
    s_var: str = "s",
    a_var: str = "a",
    post_var: str = "post",
) -> str:
    """
    Return an IML snippet that binds `post_var` to the result of applying
    the transitions in `then_names` sequentially.

    Output shape:
      - if []:    let post : state = s
      - if [t]:   let post : state = (t s a)
      - if [t1;t2]:
          let s1 : state = (t1 s a) in
          let post : state = (t2 s1 a)
    """
    ts = [t for t in (then_names or []) if str(t).strip()]
    if not ts:
        return f"let {post_var} : state = {s_var}"

    if len(ts) == 1:
        t0 = ts[0]
        return f"let {post_var} : state = ({t0} {s_var} {a_var})"

    lines: list[str] = []
    prev = s_var
    # build intermediate lets
    for i, t in enumerate(ts[:-1], start=1):
        si = f"s{i}"
        lines.append(f"let {si} : state = ({t} {prev} {a_var}) in")
        prev = si
    # final step binds post
    last = ts[-1]
    lines.append(f"let {post_var} : state = ({last} {prev} {a_var})")
    return "\n".join(lines)