import { z } from "zod";
import { openai, SCRIPT_MODEL } from "../openai";
import { buildScriptMessages, buildSectionMessages } from "../prompts/script";
import type {
  EpisodeConfig,
  ScriptProvider,
  ScriptSection,
  StructuredScript,
  TokenUsage,
} from "../types";

const turnSchema = z.object({
  speakerKey: z.string().min(1),
  text: z.string().min(1),
});

const sectionSchema = z.object({
  id: z.string().min(1),
  title: z.string().min(1),
  turns: z.array(turnSchema).min(1),
});

const scriptSchema = z.object({
  title: z.string().min(1),
  sections: z.array(sectionSchema).min(1),
});

/** Coerce/repair speakerKeys the model may have invented to the configured set. */
function normalizeSpeakers(script: StructuredScript, config: EpisodeConfig): StructuredScript {
  const valid = new Set(config.speakers.map((s) => s.speakerKey));
  const fallback = config.speakers[0]?.speakerKey ?? "host";
  return {
    ...script,
    sections: script.sections.map((sec) => ({
      ...sec,
      turns: sec.turns.map((t) => ({
        ...t,
        speakerKey: valid.has(t.speakerKey) ? t.speakerKey : fallback,
      })),
    })),
  };
}

function usageFrom(u: { prompt_tokens?: number; completion_tokens?: number } | undefined): TokenUsage {
  return { inputTokens: u?.prompt_tokens ?? 0, outputTokens: u?.completion_tokens ?? 0 };
}

export class OpenAIScriptProvider implements ScriptProvider {
  readonly model = SCRIPT_MODEL;

  async generate(config: EpisodeConfig): Promise<{ script: StructuredScript; usage: TokenUsage }> {
    const res = await openai().chat.completions.create({
      model: this.model,
      messages: buildScriptMessages(config),
      response_format: { type: "json_object" },
      temperature: 0.8,
    });
    const content = res.choices[0]?.message?.content ?? "{}";
    const parsed = scriptSchema.parse(JSON.parse(content));
    return { script: normalizeSpeakers(parsed, config), usage: usageFrom(res.usage) };
  }

  async regenerateSection(
    config: EpisodeConfig,
    script: StructuredScript,
    sectionId: string
  ): Promise<{ section: ScriptSection; usage: TokenUsage }> {
    const res = await openai().chat.completions.create({
      model: this.model,
      messages: buildSectionMessages(config, script, sectionId),
      response_format: { type: "json_object" },
      temperature: 0.9,
    });
    const content = res.choices[0]?.message?.content ?? "{}";
    const section = sectionSchema.parse(JSON.parse(content));
    const valid = new Set(config.speakers.map((s) => s.speakerKey));
    const fallback = config.speakers[0]?.speakerKey ?? "host";
    return {
      section: {
        ...section,
        id: sectionId,
        turns: section.turns.map((t) => ({
          ...t,
          speakerKey: valid.has(t.speakerKey) ? t.speakerKey : fallback,
        })),
      },
      usage: usageFrom(res.usage),
    };
  }
}