import { tool, Tool } from "ai";
import { z } from "zod";
import _ from "lodash";
import ResTool from "@/socket/resTool";
import u from "@/utils";

const deriveAssetSchema = z.object({
  id: z.number().describe("衍生资产ID,如果新增则为空"),
  assetsId: z.number().describe("关联的资产ID"),
  prompt: z.string().describe("生成提示词"),
  name: z.string().describe("衍生资产名称"),
  desc: z.string().describe("衍生资产描述"),
  src: z.string().nullable().describe("衍生资产资源路径"),
  state: z.enum(["未生成", "生成中", "已完成", "生成失败"]).describe("衍生资产生成状态"),
  type: z.enum(["role", "tool", "scene", "clip"]).describe("衍生资产类型"),
});
const assetItemSchema = z.object({
  id: z.number().describe("资产唯一标识"),
  name: z.string().describe("资产名称"),
  type: z.enum(["role", "tool", "scene", "clip"]).describe("资产类型"),
  prompt: z.string().describe("生成提示词"),
  desc: z.string().describe("资产描述"),
  derive: z.array(deriveAssetSchema).describe("衍生资产列表"),
});
const storyboardSchema = z.object({
  id: z.number().describe("分镜ID，必须为真实id"),
  title: z.string().describe("分镜标题"),
  description: z.string().describe("分镜描述"),
  camera: z.string().describe("镜头信息"),
  duration: z.number().describe("持续时长(秒)"),
  frameMode: z.enum(["firstFrame", "endFrame", "linesSoundEffects"]).describe("帧模式: 首帧/尾帧/台词音效"),
  prompt: z.string().describe("生成提示词"),
  lines: z.string().nullable().describe("台词内容"),
  sound: z.string().nullable().describe("音效内容"),
  mode: z
    .union([
      z.enum(["singleImage", "multiImage", "gridImage", "startEndRequired", "endFrameOptional", "startFrameOptional", "text"]),
      z.array(z.enum(["video", "image", "audio", "text"])),
    ])
    .describe("视频模式"),
  associateAssetsIds: z.array(z.number()).describe("关联资产ID列表"),
  src: z.string().nullable().describe("分镜资源路径"),
});
const workbenchDataSchema = z.object({
  name: z.string().describe("项目名称"),
  duration: z.string().describe("视频时长"),
  resolution: z.string().describe("分辨率"),
  fps: z.string().describe("帧率"),
  cover: z.string().optional().describe("封面图片路径"),
  gradient: z.string().optional().describe("渐变色配置"),
});
const posterItemSchema = z.object({
  id: z.number().describe("海报ID"),
  image: z.string().describe("海报图片路径"),
});
const flowDataSchema = z.object({
  script: z.string().describe("剧本内容"),
  scriptPlan: z.string().describe("拍摄计划"),
  assets: z.array(assetItemSchema).describe("衍生资产"),
  storyboardTable: z.string().describe("分镜表"),
  storyboard: z.array(storyboardSchema).describe("分镜面板"),
});

export type FlowData = z.infer<typeof flowDataSchema>;

const keySchema = z.enum(Object.keys(flowDataSchema.shape) as [keyof FlowData, ...Array<keyof FlowData>]);
const flowDataKeyLabels = Object.fromEntries(
  Object.entries(flowDataSchema.shape).map(([key, schema]) => [key, (schema as z.ZodTypeAny).description ?? key]),
) as Record<keyof FlowData, string>;

interface ToolConfig {
  resTool: ResTool;
  toolsNames?: string[];
  msg: ReturnType<ResTool["newMessage"]>;
}

export default (toolCpnfig: ToolConfig) => {
  const { resTool, toolsNames, msg } = toolCpnfig;
  const { socket } = resTool;
  const tools: Record<string, Tool> = {
    get_flowData: tool({
      description: "获取工作区数据",
      inputSchema: z.object({
        key: keySchema.describe("数据key"),
      }),
      execute: async ({ key }) => {
        const thinking = msg.thinking(`正在获取${flowDataKeyLabels[key]}工作区数据...`);
        console.log("[tools] get_flowData", key);
        const flowData: FlowData = await new Promise((resolve) => socket.emit("getFlowData", { key }, (res: any) => resolve(res)));
        thinking.appendText(`获取到${flowDataKeyLabels[key]}:\n` + JSON.stringify(flowData[key], null, 2));
        thinking.updateTitle(`获取${flowDataKeyLabels[key]}完成`);
        thinking.complete();
        return flowData[key];
      },
    }),
    add_deriveAsset: tool({
      description: "新增或更新衍生资产",
      inputSchema: z.object({
        assetsId: z.number().describe("关联的资产ID"),
        id: z.number().nullable().describe("衍生资产ID,如果新增则为空"),
        name: z.string().describe("衍生资产名称"),
        desc: z.string().describe("衍生资产描述"),
        type: z.enum(["role", "tool", "scene", "clip"]).describe("衍生资产类型"),
      }),
      execute: async (deriveAsset) => {
        const thinking = msg.thinking("正在操作资产...");
        const { projectId, scriptId } = resTool.data;
        const startTime = Date.now();
        const data = {
          id: deriveAsset.id ?? undefined,
          assetsId: deriveAsset.assetsId,
          projectId,
          name: deriveAsset.name,
          type: deriveAsset.type,
          describe: deriveAsset.desc,
          startTime,
        };
        if (deriveAsset.id) {
          await u.db("o_assets").where("id", deriveAsset.id).update(data);
          thinking.appendText(`已更新衍生资产，ID: ${deriveAsset.id}\n`);
        } else {
          const [insertedId] = await u.db("o_assets").insert(data);
          data.id = insertedId;
          await u.db("o_scriptAssets").insert({ scriptId, assetId: insertedId });
          thinking.appendText(`已新增衍生资产，ID: ${insertedId}\n`);
        }
        const res = await new Promise((resolve) => socket.emit("addDeriveAsset", data, (res: any) => resolve(res)));
        thinking.updateTitle("资产操作完成");
        thinking.complete();
        return res ?? "操作成功";
      },
    }),
    del_deriveAsset: tool({
      description: "删除衍生资产",
      inputSchema: z.object({
        assetsId: z.number().describe("关联的资产ID"),
        id: z.number().describe("衍生资产ID"),
      }),
      execute: async ({ assetsId, id }) => {
        const thinking = msg.thinking("正在操作资产...");
        const { scriptId } = resTool.data;
        await u.db("o_assets").where("id", id).del();
        await u.db("o_scriptAssets").where({ scriptId, assetId: id }).del();
        thinking.appendText(`已删除衍生资产，ID: ${id}\n`);
        const res = await new Promise((resolve) => socket.emit("delDeriveAsset", { assetsId, id }, (res: any) => resolve(res)));
        thinking.updateTitle("资产操作完成");
        thinking.complete();
        return res ?? "删除成功";
      },
    }),

    add_storyboard: tool({
      description: "新增或更新分镜面板",
      inputSchema: z.object({
        id: z.number().nullable().describe("分镜面板ID,如果新增则为空"),
        title: z.string().describe("分镜面板名称"),
        desc: z.string().describe("分镜面板描述"),
        group: z.number().describe("分镜面板分组,根据这个字段 对分镜图片，进行同时生成视频，例如 同一分组的两张图片会被用于首尾帧生成视频"),
      }),
      execute: async (storyboard) => {
        const thinking = msg.thinking("正在操作资产...");
        const { projectId, scriptId } = resTool.data;
        const createTime = Date.now();
        console.log("%c Line:161 🍤 storyboard", "background:#e41a6a", storyboard);

        const data = {
          id: storyboard.id ?? undefined,
          title: storyboard.title,
          description: storyboard.desc,
          createTime,
          scriptId,
        };
        if (storyboard.id) {
          await u.db("o_storyboard").where("id", storyboard.id).update(data);
          thinking.appendText(`已更新分镜面板，ID: ${storyboard.id}\n`);
        } else {
          const [insertedId] = await u.db("o_storyboard").insert(data);
          data.id = insertedId;
          thinking.appendText(`已新增分镜面板，ID: ${insertedId}\n`);
        }
        const res = await new Promise((resolve) => socket.emit("addStoryboard", data, (res: any) => resolve(res)));
        thinking.updateTitle("分镜面板操作完成");
        thinking.complete();
        return res ?? "操作成功";
      },
    }),
    generate_deriveAsset: tool({
      description: "生成衍生资产",
      inputSchema: z.object({
        id: z.array(z.number()).describe("需要生成的 衍生资产ID"),
      }),
      execute: async ({ id }) => {
        const thinking = msg.thinking("正在生成衍生资产...");
        const res = await new Promise((resolve) => socket.emit("generateDeriveAsset", { id }, (res: any) => resolve(res)));
        thinking.appendText(`已生成衍生资产，ID: ${id}\n`);
        thinking.updateTitle("衍生资产生成完成");
        thinking.complete();
        return res ?? "生成失败";
      },
    }),
    generate_storyboard: tool({
      description: "生成分镜图片",
      inputSchema: z.object({
        storyboardIds: z.array(z.number()).describe("分镜ID列表"),
      }),
      execute: async ({ storyboardIds }) => {
        const thinking = msg.thinking("正在生成分镜...");
        const res = await new Promise((resolve) => socket.emit("generateStoryboard", { storyboardIds }, (res: any) => resolve(res)));
        thinking.appendText("生成的分镜数据:\n" + JSON.stringify(res, null, 2));
        thinking.updateTitle("分镜生成完成");
        thinking.complete();
        return res;
      },
    }),
  };

  return toolsNames ? Object.fromEntries(Object.entries(tools).filter(([n]) => toolsNames.includes(n))) : tools;
};