import type {
  ColineAiEvent,
  ColineAiJson,
  ColineAiJob,
  ColineAiModelInfo,
  ColineAiRequest,
  ColineAi,
  ColineImageOptions,
} from "../capabilities/ai";
interface Fixture {
  events: ColineAiEvent[];
  job: ColineAiJob;
  owner: string;
  key: string;
  request: string;
}
const fixtures = new WeakMap<object, Map<string, Fixture>>();
export function fakeAi(
  state: object,
  owner: string | null,
  models: ColineAiModelInfo[],
  handler?: (
    request: ColineAiRequest | { kind: "image"; input: ColineImageOptions },
    context: { event: (type: string, data: ColineAiJson) => Promise<void> },
  ) => Promise<ColineAiJob["output"]>,
): ColineAi {
  let jobs = fixtures.get(state);
  if (!jobs) {
    jobs = new Map();
    fixtures.set(state, jobs);
  }
  const records = jobs;
  const get = (id: string) => {
    const value = records.get(id);
    if (!owner || value?.owner !== owner) throw new Error("job_not_found");
    return structuredClone(value.job);
  };
  const run = async (request: ColineAiRequest | { kind: "image"; input: ColineImageOptions }) => {
    if (!owner) throw new Error("user_required");
    const encoded = JSON.stringify(request);
    const prior = [...records.values()].find(
      (value) => value.owner === owner && value.key === request.input.idempotencyKey,
    );
    if (prior) {
      if (prior.request !== encoded) throw new Error("idempotency_conflict");
      return structuredClone(prior.job);
    }
    if (!handler)
      throw new Error("Provide onAiJob to test AI jobs. Tests never call real providers.");
    const model = models.find(
      (value) => value.id === request.input.model && value.kind === request.kind,
    );
    if (!model) throw new Error("ai_model_unavailable");
    const count =
      request.kind === "image"
        ? (request.input.count ?? 1)
        : request.kind === "language"
          ? (request.input.maxSteps ?? 1)
          : request.kind === "video"
            ? (request.input.count ?? 1)
            : 1;
    const price = model.priceMicrousd * count;
    if (price > request.input.maxCostMicrousd) throw new Error("ai_budget_exceeded");
    const job: ColineAiJob = {
      id: `ai_${records.size + 1}`,
      kind: request.kind,
      status: "queued",
      model: model.id,
      images: [],
      previews: [],
      error: null,
      chargeMicrousd: price,
      createdAt: new Date().toISOString(),
      updatedAt: new Date().toISOString(),
    };
    records.set(job.id, {
      events: [],
      job,
      owner,
      key: request.input.idempotencyKey,
      request: encoded,
    });
    const initial = structuredClone(job);
    void Promise.resolve().then(async () => {
      if (job.status !== "queued") return;
      job.status = "running";
      try {
        const output = await handler(request, {
          event: async (type, data) => {
            if (job.status !== "running") throw new Error("cancelled");
            const events = records.get(job.id)!.events;
            events.push({ sequence: events.length + 1, type, data: structuredClone(data) });
          },
        });
        if (job.status === "running") {
          if (output !== undefined) job.output = output;
          if (
            request.kind === "image" &&
            output &&
            typeof output === "object" &&
            !Array.isArray(output) &&
            Array.isArray(output.images)
          ) {
            job.images = output.images.map((asset) => {
              if (
                !asset ||
                typeof asset !== "object" ||
                Array.isArray(asset) ||
                typeof asset.fileId !== "string" ||
                typeof asset.mimeType !== "string"
              )
                throw new Error("Invalid fake image asset");
              return { fileId: asset.fileId, mimeType: asset.mimeType };
            });
          }
          if (
            request.kind === "language" &&
            output &&
            typeof output === "object" &&
            !Array.isArray(output) &&
            typeof output.stepCount === "number"
          ) {
            if (
              !Number.isInteger(output.stepCount) ||
              output.stepCount < 1 ||
              output.stepCount > count
            )
              throw new Error("Invalid step count");
            job.chargeMicrousd = model.priceMicrousd * output.stepCount;
          }
          job.status = "succeeded";
        }
      } catch {
        if (job.status === "running") {
          job.status = "failed";
          job.chargeMicrousd = 0;
          job.error = { code: "ai_generation_failed", message: "Test provider failed." };
        }
      }
      job.updatedAt = new Date().toISOString();
    });
    return initial;
  };
  return {
    models: async () => structuredClone(models),
    run,
    images: { generate: (input) => run({ kind: "image", input }) },
    jobs: {
      get: async (id) => get(id),
      events: async (id, afterSequence = 0) => {
        get(id);
        if (!Number.isInteger(afterSequence) || afterSequence < 0)
          throw new Error("Invalid cursor");
        return structuredClone(
          records
            .get(id)!
            .events.filter((event) => event.sequence > afterSequence)
            .slice(0, 200),
        );
      },
      list: async (input) =>
        [...records.values()]
          .filter((value) => value.owner === owner)
          .slice(-(input?.limit ?? 30))
          .reverse()
          .map((value) => {
            const summary = structuredClone(value.job);
            delete summary.output;
            return summary;
          }),
      cancel: async (id) => {
        get(id);
        const job = records.get(id)!.job;
        if (["queued", "running"].includes(job.status)) {
          job.status = "cancelled";
          job.chargeMicrousd = 0;
        }
        return structuredClone(job);
      },
    },
  };
}