import type { EffectBackend } from "../../../../src/core/effects/backend";
import type { State } from "../../../../src/core/eval/machine";
import type { OpCall, Resumption } from "../../../../src/core/effects/opcall";
import type { Val } from "../../../../src/core/eval/values";
import type { BudgetTracker } from "../../../../src/package-bridge/governance";

type AmbAlternative = {
  thunk: Val;
  resumption: Resumption;
};

export type AmbStrategy = "DFS" | "FAIR";

export class AmbBackend implements EffectBackend {
  readonly ops = ["amb.choose", "amb.fail"] as const;

  private alternatives: AmbAlternative[] = [];
  private strategy: AmbStrategy = "DFS";

  constructor(private budget?: BudgetTracker) {}

  setStrategy(strategy: AmbStrategy): void {
    this.strategy = strategy;
  }

  reset(): void {
    this.alternatives = [];
  }

  hasAlternatives(): boolean {
    return this.alternatives.length > 0;
  }

  setBudget(budget: BudgetTracker): void {
    this.budget = budget;
  }

  async handle(state: State, opcall: OpCall): Promise<State> {
    if (opcall.op === "amb.choose") {
      return this.handleChoose(state, opcall);
    }
    if (opcall.op === "amb.fail") {
      return this.handleFail(state, opcall);
    }
    throw new Error(`AmbBackend: unhandled op ${opcall.op}`);
  }

  private handleChoose(state: State, opcall: OpCall): State {
    const thunks = opcall.args[0];
    if (!thunks || thunks.tag !== "Vector") {
      throw new Error("amb.choose: expected vector of thunks");
    }

    function consToArray(v: Val): Val[] {
      const result: Val[] = [];
      let cur = v;
      while (cur.tag === "Vector" && cur.items.length === 2) {
        result.push(cur.items[0]);
        cur = cur.items[1];
      }
      if (cur.tag === "Vector" && cur.items.length > 2) {
        return cur.items;
      }
      return result;
    }

    const items = consToArray(thunks);
    if (items.length === 0) {
      return this.handleFail(state, {
        op: "amb.fail",
        args: [{ tag: "Str", s: "amb: no alternatives" } as Val],
        resumption: opcall.resumption,
        ctxDigest: opcall.ctxDigest,
      });
    }

    for (let i = 1; i < items.length; i++) {
      this.alternatives.push({
        thunk: items[i],
        resumption: opcall.resumption,
      });
    }

    this.budget?.consumeAmbAttempt?.();

    const firstThunk = items[0];
    if (firstThunk.tag !== "Closure") {
      throw new Error("amb.choose: thunk must be a closure");
    }

    const baseState: State | "Uncaught" | "OutOfBudget" =
      opcall.resumption.invoke({ tag: "Unit" } as Val);
    if (baseState === ("Uncaught" as any) || baseState === ("OutOfBudget" as any)) {
      return baseState as any;
    }

    return {
      ...baseState,
      control: { tag: "Expr", e: firstThunk.body },
      env: firstThunk.env,
    } as State;
  }

  private handleFail(_state: State, _opcall: OpCall): State {
    if (this.alternatives.length === 0) {
      throw new Error("amb: all alternatives exhausted");
    }

    this.budget?.consumeAmbAttempt?.();

    const alt = this.alternatives.shift()!;
    const thunk = alt.thunk;
    const resumption = alt.resumption;

    const baseState: State | "Uncaught" | "OutOfBudget" =
      resumption.invoke({ tag: "Unit" } as Val);
    if (baseState === ("Uncaught" as any) || baseState === ("OutOfBudget" as any)) {
      return baseState as any;
    }

    if (thunk.tag === "Closure") {
      return {
        ...baseState,
        control: { tag: "Expr", e: thunk.body },
        env: thunk.env,
      } as State;
    }
    if (thunk.tag === "Native") {
      throw new Error("amb: Native thunks not yet supported");
    }

    throw new Error(`amb: invalid thunk type ${thunk.tag}`);
  }
}