import { randomUUID } from "node:crypto";
import type { EventStore } from "./event-store.js";
import { cloneMessages, firstNewUserMessage, stableHash } from "./text.js";
import type {
  BranchBackend,
  CandidateRecord,
  ForkResult,
  PreferenceRecord,
  RunResult,
  SerializedMessage,
  Side,
  StreamUpdate,
} from "./types.js";

export interface ForkCallbacks {
  left?: (update: StreamUpdate) => void;
  right?: (update: StreamUpdate) => void;
}

export class BranchEngine {
  private history: SerializedMessage[] = [];

  constructor(
    readonly backend: BranchBackend,
    readonly store: EventStore,
  ) {}

  get canonicalHistory(): SerializedMessage[] {
    return cloneMessages(this.history);
  }

  restore(messages: SerializedMessage[]): void {
    this.history = cloneMessages(messages);
  }

  async fork(prompt: string, callbacks: ForkCallbacks = {}): Promise<ForkResult> {
    const turnId = randomUUID();
    const parentHistory = cloneMessages(this.history);
    const [left, right] = await Promise.all([
      this.backend.run({
        history: cloneMessages(parentHistory),
        prompt,
        branchId: `${turnId}:left`,
        ...(callbacks.left ? { onUpdate: callbacks.left } : {}),
      }),
      this.backend.run({
        history: cloneMessages(parentHistory),
        prompt,
        branchId: `${turnId}:right`,
        ...(callbacks.right ? { onUpdate: callbacks.right } : {}),
      }),
    ]);

    return {
      turnId,
      createdAt: new Date().toISOString(),
      prompt,
      parentHistory,
      parentHistoryHash: stableHash(parentHistory),
      left,
      right,
    };
  }

  async select(fork: ForkResult, side: Side): Promise<PreferenceRecord> {
    const preferred = fork[side];
    const rejected = fork[side === "left" ? "right" : "left"];
    this.history = cloneMessages(preferred.messages);

    const record = this.baseRecord(fork, [this.candidate(fork.left, "left"), this.candidate(fork.right, "right")], {
      preferredCandidateId: preferred.branchId,
      rejectedCandidateIds: [rejected.branchId],
      selection: side,
    });
    await this.store.append(record);
    return record;
  }

  async revise(
    fork: ForkResult,
    side: Side,
    feedback: string,
    onUpdate?: (update: StreamUpdate) => void,
  ): Promise<PreferenceRecord> {
    const original = fork[side];
    const other = fork[side === "left" ? "right" : "left"];
    const revisionPrompt = buildRevisionPrompt(fork.prompt, original.response, feedback);
    const revision = await this.backend.run({
      history: cloneMessages(original.messages),
      prompt: revisionPrompt,
      branchId: `${fork.turnId}:revision`,
      ...(onUpdate ? { onUpdate } : {}),
    });

    const originalUserMessage = firstNewUserMessage(fork.parentHistory.length, original.messages);
    const revisedAssistantMessage = revision.messages.findLast((message) => message.role === "assistant");
    if (!originalUserMessage || !revisedAssistantMessage) {
      throw new Error("Could not squash the correction into a clean canonical turn.");
    }

    this.history = [
      ...cloneMessages(fork.parentHistory),
      structuredClone(originalUserMessage),
      structuredClone(revisedAssistantMessage),
    ];

    const candidates = [
      this.candidate(fork.left, "left"),
      this.candidate(fork.right, "right"),
      this.candidate(revision, "revision"),
    ];
    const record = this.baseRecord(fork, candidates, {
      preferredCandidateId: revision.branchId,
      rejectedCandidateIds: [original.branchId, other.branchId],
      selection: side === "left" ? "revise-left" : "revise-right",
      correction: {
        feedback,
        originalCandidateId: original.branchId,
        revisionPrompt,
      },
    });
    await this.store.append(record);
    return record;
  }

  private candidate(result: RunResult, source: CandidateRecord["source"]): CandidateRecord {
    return {
      id: result.branchId,
      source,
      response: result.response,
      reasoning: result.reasoning,
      messages: cloneMessages(result.messages),
      model: structuredClone(result.model),
    };
  }

  private baseRecord(
    fork: ForkResult,
    candidates: CandidateRecord[],
    decision: Pick<
      PreferenceRecord,
      "preferredCandidateId" | "rejectedCandidateIds" | "selection" | "correction"
    >,
  ): PreferenceRecord {
    return {
      schema: "branch-preference/v1",
      id: fork.turnId,
      createdAt: new Date().toISOString(),
      backend: this.backend.name,
      systemPrompt: this.backend.systemPrompt,
      systemPromptHash: stableHash(this.backend.systemPrompt),
      parentHistory: cloneMessages(fork.parentHistory),
      parentHistoryHash: fork.parentHistoryHash,
      prompt: fork.prompt,
      candidates,
      preferredCandidateId: decision.preferredCandidateId,
      rejectedCandidateIds: decision.rejectedCandidateIds,
      selection: decision.selection,
      ...(decision.correction ? { correction: decision.correction } : {}),
    };
  }
}

export function buildRevisionPrompt(originalPrompt: string, originalResponse: string, feedback: string): string {
  return [
    "The response you just gave violated an important preference or instruction.",
    "",
    "Feedback:",
    feedback,
    "",
    "Original user prompt:",
    originalPrompt,
    "",
    "Response being corrected:",
    originalResponse,
    "",
    "Respond to the original user prompt again as if you had already understood this feedback.",
    "Do not mention the correction, this feedback, or a previous response. Return only the improved response.",
  ].join("\n");
}