import type { AgentRuntime, AgentRuntimeEvent } from "@rakazo/adapter-kit";
import type { PrismaClient } from "@rakazo/db";
import { describe, expect, it, vi } from "vitest";
import { z } from "zod";
import { createMediaModelClient, extractJson } from "./media-context.js";

function runtime(reply: string, compaction = true): AgentRuntime {
  return {
    describe: () => ({ capabilities: { compaction } }),
    run: vi.fn(async function* (): AsyncGenerator<AgentRuntimeEvent> {
      yield { type: "done", text: reply };
    }),
  } as unknown as AgentRuntime;
}

function prisma(consented: boolean): PrismaClient {
  return {
    aiDataConsent: { findUnique: vi.fn(async () => (consented ? { userId: "u" } : null)) },
  } as unknown as PrismaClient;
}

const model = async () => ({ provider: "openai", id: "gpt-test" });
const input = {
  spaceId: "s",
  userId: "u",
  instructions: "Reply with JSON.",
  prompt: "Label",
  schema: z.object({ title: z.string() }),
};

describe("extractJson", () => {
  it("reads fenced, wrapped, and bare JSON", () => {
    expect(extractJson('```json\n{"a":1}\n```')).toEqual({ a: 1 });
    expect(extractJson('Thinking... {"a":[1,2]} done')).toEqual({ a: [1, 2] });
    expect(extractJson("[1]")).toEqual([1]);
    expect(extractJson("no json")).toBeUndefined();
    expect(extractJson("{broken")).toBeUndefined();
  });
});

describe("createMediaModelClient", () => {
  it("returns parsed output only with consent for the resolved recipient", async () => {
    const withConsent = createMediaModelClient({
      prisma: prisma(true),
      runtime: runtime('{"title":"Beach"}'),
      resolveModel: model,
    });
    expect(await withConsent.generate(input)).toEqual({ title: "Beach" });

    const run = runtime('{"title":"Beach"}');
    const withoutConsent = createMediaModelClient({
      prisma: prisma(false),
      runtime: run,
      resolveModel: model,
    });
    expect(await withoutConsent.generate(input)).toBeNull();
    expect(run.run).not.toHaveBeenCalled();
  });

  it("falls back to null for scripted runtimes, missing models, and invalid replies", async () => {
    const scripted = createMediaModelClient({
      prisma: prisma(true),
      runtime: runtime('{"title":"x"}', false),
      resolveModel: model,
    });
    expect(await scripted.generate(input)).toBeNull();
    const noModel = createMediaModelClient({
      prisma: prisma(true),
      runtime: runtime('{"title":"x"}'),
      resolveModel: async () => {
        throw new Error("no model");
      },
    });
    expect(await noModel.generate(input)).toBeNull();
    const invalid = createMediaModelClient({
      prisma: prisma(true),
      runtime: runtime('{"name":"x"}'),
      resolveModel: model,
    });
    expect(await invalid.generate(input)).toBeNull();
  });
});