import { v } from "convex/values";
import OpenAI from "openai";
import { createSunoClient } from "../lib/sunoClient";
import { createFalClient } from "../lib/falClient";
import { internal } from "./_generated/api";
import { Id } from "./_generated/dataModel";
import {
  internalAction,
  internalMutation,
  internalQuery,
  mutation,
  query,
} from "./_generated/server";
import { getAuthUserId } from "@convex-dev/auth/server";

// Query atoms by IDs
export const getByIds = query({
  args: {
    ids: v.array(v.id("atoms")),
  },
  handler: async (ctx, args) => {
    const userId = await getAuthUserId(ctx);
    if (!userId) throw new Error("Not authenticated");

    return await Promise.all(args.ids.map((id) => ctx.db.get(id)));
  },
});

// Query atoms by space
export const getBySpace = query({
  args: {
    spaceId: v.id("spaces"),
    type: v.optional(
      v.union(
        v.literal("song"),
        v.literal("video"),
        v.literal("image"),
        v.literal("lyrics"),
        v.literal("webview"),
        v.literal("midi")
      )
    ),
  },
  handler: async (ctx, args) => {
    const userId = await getAuthUserId(ctx);
    if (!userId) throw new Error("Not authenticated");

    // Check if user is a member of the space
    const member = await ctx.db
      .query("spaceMembers")
      .withIndex("by_space_and_user", (q) =>
        q.eq("spaceId", args.spaceId).eq("userId", userId)
      )
      .first();

    if (!member || member.isKicked) {
      throw new Error("Not authorized - not a member of this space");
    }

    if (args.type !== undefined) {
      return await ctx.db
        .query("atoms")
        .withIndex("by_space_and_type", (q) =>
          q.eq("spaceId", args.spaceId).eq("type", args.type!)
        )
        .collect();
    }
    return await ctx.db
      .query("atoms")
      .withIndex("by_space", (q) => q.eq("spaceId", args.spaceId))
      .collect();
  },
});

// Create atom (internal only)
export const createInternal = internalMutation({
  args: {
    type: v.union(
      v.literal("song"),
      v.literal("video"),
      v.literal("image"),
      v.literal("lyrics"),
      v.literal("webview"),
      v.literal("midi")
    ),
    spaceId: v.id("spaces"),
    ownerId: v.id("users"),
    metadata: v.any(),
    status: v.union(
      v.literal("pending"),
      v.literal("processing"),
      v.literal("streaming"),
      v.literal("completed"),
      v.literal("failed")
    ),
    progress: v.number(),
  },
  handler: async (ctx, args) => {
    return await ctx.db.insert("atoms", {
      ...args,
      createdAt: Date.now(),
      updatedAt: Date.now(),
    });
  },
});

// Update atom metadata and status (internal only)
export const updateInternal = internalMutation({
  args: {
    id: v.id("atoms"),
    metadata: v.optional(v.any()),
    status: v.optional(
      v.union(
        v.literal("pending"),
        v.literal("processing"),
        v.literal("streaming"),
        v.literal("completed"),
        v.literal("failed")
      )
    ),
    progress: v.optional(v.number()),
  },
  handler: async (ctx, args) => {
    const { id, ...updates } = args;
    await ctx.db.patch(id, {
      ...updates,
      updatedAt: Date.now(),
    });
  },
});

// Generate song action
export const generateSong = internalAction({
  args: {
    lyrics: v.optional(v.string()),
    tags: v.optional(v.string()),
    makeInstrumental: v.optional(v.boolean()),
    coverClipId: v.optional(v.string()),
    spaceId: v.id("spaces"),
    userId: v.id("users"),
    radioPrompt: v.optional(v.string()), // Store original radio prompt for queue management
  },
  handler: async (ctx, args): Promise<Array<any>> => {
    // Get user's Suno token if they have one
    const userToken = await ctx.runQuery(internal.users.getUserTokenInternal, {
      userId: args.userId,
    });

    const client = createSunoClient(userToken);

    try {
      // 1. Generate songs via Suno API (returns 2 song IDs)
      // Pass lyrics as prompt to Suno API
      const songIds = await client.generateSongs({
        prompt: args.lyrics ?? "",
        tags: args.tags,
        makeInstrumental: args.makeInstrumental ?? false,
        coverClipId: args.coverClipId,
        task: args.coverClipId ? "cover" : undefined,
      });

      // 2. Create atom records for both songs
      const atomIds: Array<any> = await Promise.all(
        songIds.map((songId: string) =>
          ctx.runMutation(internal.atoms.createInternal, {
            type: "song",
            spaceId: args.spaceId,
            ownerId: args.userId,
            metadata: {
              sunoClipId: songId,
              title: "Generating...",
              lyrics: args.lyrics,
              tags: args.tags,
              coverClipId: args.coverClipId,
              radioPrompt: args.radioPrompt, // Store radio prompt for filtering
            },
            status: "pending",
            progress: 0,
          })
        )
      );

      // 3. Schedule polling
      await ctx.scheduler.runAfter(0, internal.atoms.pollSongStatus, {
        atomIds,
      });

      return atomIds;
    } catch (error) {
      console.error("Song generation failed:", error);
      throw error;
    }
  },
});

// Poll song status (internal scheduled function)
export const pollSongStatus = internalAction({
  args: {
    atomIds: v.array(v.id("atoms")),
  },
  handler: async (ctx, args) => {
    // Get atoms
    const atoms = await Promise.all(
      args.atomIds.map((id) =>
        ctx.runQuery(internal.atoms.getAtomInternal, { id })
      )
    );

    // Filter active atoms
    const activeAtoms = atoms.filter(
      (atom: any) => atom && !["completed", "failed"].includes(atom.status)
    );

    if (activeAtoms.length === 0) return;

    // Get user token from first atom's owner (all atoms in a batch share the same owner)
    const userId = activeAtoms[0]?.ownerId;
    const userToken = userId
      ? await ctx.runQuery(internal.users.getUserTokenInternal, { userId })
      : null;

    // Poll Suno API
    const client = createSunoClient(userToken);
    const sunoIds = activeAtoms.map((a: any) => a.metadata.sunoClipId);

    try {
      const songs = await client.getSongStatus(sunoIds);

      // Update atoms
      for (let i = 0; i < activeAtoms.length; i++) {
        const atom = activeAtoms[i];
        if (!atom) continue;
        const song = songs[i];

        const hadNoDuration = !atom.metadata?.duration;
        const nowHasDuration = !!song.duration;

        await ctx.runMutation(internal.atoms.updateInternal, {
          id: atom._id,
          metadata: {
            ...atom.metadata,
            title: song.title || atom.metadata.title,
            artist: song.artist,
            audioUrl: song.audioUrl,
            videoUrl: song.videoUrl,
            albumArtUrl: song.albumArtUrl,
            duration: song.duration,
            sunoStatus: song.status,
          },
          status: song.status,
          progress:
            song.status === "completed"
              ? 100
              : song.status === "streaming"
                ? 75
                : song.status === "processing"
                  ? 50
                  : 25,
        });

        // If duration just became available, update radio room schedules
        if (hadNoDuration && nowHasDuration && song.duration) {
          await ctx.runMutation(internal.radioPlayback.updateTrackDuration, {
            atomId: atom._id,
            durationSeconds: song.duration,
          });
        }
      }

      // Schedule next poll if needed
      const stillActive = songs.some(
        (s) => !["completed", "failed"].includes(s.status)
      );

      if (stillActive) {
        await ctx.scheduler.runAfter(2000, internal.atoms.pollSongStatus, {
          atomIds: args.atomIds,
        });
      }
    } catch (error) {
      console.error("Polling error:", error);
      // Retry polling after delay
      await ctx.scheduler.runAfter(2000, internal.atoms.pollSongStatus, {
        atomIds: args.atomIds,
      });
    }
  },
});

// Get atom by ID (internal only)
export const getAtomInternal = internalQuery({
  args: {
    id: v.id("atoms"),
  },
  handler: async (ctx, args) => {
    return await ctx.db.get(args.id);
  },
});

// Get storage URL for an image atom
export const getImageUrl = query({
  args: {
    storageId: v.string(),
  },
  handler: async (ctx, args) => {
    const userId = await getAuthUserId(ctx);
    if (!userId) throw new Error("Not authenticated");

    return await ctx.storage.getUrl(args.storageId as any);
  },
});

// Toggle like on an atom
export const toggleLike = mutation({
  args: {
    atomId: v.id("atoms"),
  },
  handler: async (ctx, args) => {
    const userId = await getAuthUserId(ctx);
    if (!userId) {
      throw new Error("Not authenticated");
    }

    const atom = await ctx.db.get(args.atomId);
    if (!atom) {
      throw new Error("Atom not found");
    }

    const likedBy = atom.likedBy ?? [];
    const hasLiked = likedBy.includes(userId);

    const newLikedBy = hasLiked
      ? likedBy.filter((id) => id !== userId)
      : [...likedBy, userId];

    await ctx.db.patch(args.atomId, {
      likedBy: newLikedBy,
      likeCount: newLikedBy.length,
      updatedAt: Date.now(),
    });

    // If it's a song atom, also update reaction in Suno API
    if (atom.type === "song" && atom.metadata?.sunoClipId) {
      await ctx.scheduler.runAfter(
        0,
        internal.atoms.updateSunoReaction,
        {
          sunoClipId: atom.metadata.sunoClipId,
          reaction: hasLiked ? null : "LIKE", // null to remove like, "LIKE" to add
          userId,
        }
      );
    }

    return { liked: !hasLiked, likeCount: newLikedBy.length };
  },
});

// Increment play count on an atom
export const incrementPlayCount = mutation({
  args: {
    atomId: v.id("atoms"),
  },
  handler: async (ctx, args) => {
    const userId = await getAuthUserId(ctx);
    if (!userId) throw new Error("Not authenticated");

    const atom = await ctx.db.get(args.atomId);
    if (!atom) {
      throw new Error("Atom not found");
    }

    const newPlayCount = (atom.playCount ?? 0) + 1;

    await ctx.db.patch(args.atomId, {
      playCount: newPlayCount,
      updatedAt: Date.now(),
    });

    // If it's a song atom, also increment play count in Suno API
    if (atom.type === "song" && atom.metadata?.sunoClipId) {
      await ctx.scheduler.runAfter(0, internal.atoms.incrementSunoPlayCount, {
        sunoClipId: atom.metadata.sunoClipId,
        userId,
      });
    }

    return { playCount: newPlayCount };
  },
});

// Internal action to update Suno reaction (like/dislike)
export const updateSunoReaction = internalAction({
  args: {
    sunoClipId: v.string(),
    reaction: v.union(v.literal("LIKE"), v.literal("DISLIKE"), v.null()),
    userId: v.id("users"),
  },
  handler: async (ctx, args) => {
    // Get user's Suno token if they have one
    const userToken = await ctx.runQuery(internal.users.getUserTokenInternal, {
      userId: args.userId,
    });

    const client = createSunoClient(userToken);
    try {
      await client.updateReactionType(args.sunoClipId, args.reaction, {});
    } catch (error) {
      console.error("Failed to update Suno reaction:", error);
      // Don't throw - we don't want to fail the local like operation
    }
  },
});

// Internal action to increment Suno play count
export const incrementSunoPlayCount = internalAction({
  args: {
    sunoClipId: v.string(),
    userId: v.id("users"),
  },
  handler: async (ctx, args) => {
    // Get user's Suno token if they have one
    const userToken = await ctx.runQuery(internal.users.getUserTokenInternal, {
      userId: args.userId,
    });

    const client = createSunoClient(userToken);
    try {
      await client.incrementPlayCount(args.sunoClipId);
    } catch (error) {
      console.error("Failed to increment Suno play count:", error);
      // Don't throw - we don't want to fail the local play operation
    }
  },
});

// Generate images action
export const generateImage = internalAction({
  args: {
    prompt: v.string(),
    spaceId: v.id("spaces"),
    userId: v.id("users"),
  },
  handler: async (ctx, args): Promise<Array<any>> => {
    const client = createFalClient();

    try {
      // 1. Generate images via fal.ai API (returns 2 image URLs)
      const result = await client.generateImages(args.prompt, {
        numImages: 2,
      });

      // 2. Download images and upload to Convex storage
      const atomIds: Array<any> = [];

      for (let i = 0; i < result.images.length; i++) {
        const image = result.images[i];

        // Download image
        const imageBuffer = await client.downloadImage(image.url);

        // Upload to Convex storage
        const storageId = await ctx.storage.store(
          new Blob([imageBuffer], { type: image.content_type })
        );

        // Create atom record
        const atomId = await ctx.runMutation(internal.atoms.createInternal, {
          type: "image",
          spaceId: args.spaceId,
          ownerId: args.userId,
          metadata: {
            prompt: args.prompt,
            width: image.width,
            height: image.height,
            storageId,
            originalUrl: image.url,
          },
          status: "completed",
          progress: 100,
        });

        atomIds.push(atomId);
      }

      return atomIds;
    } catch (error) {
      console.error("Image generation failed:", error);
      throw error;
    }
  },
});

// Query atoms created by the current user with optional filters
export const getMyAtoms = query({
  args: {
    type: v.optional(
      v.union(
        v.literal("song"),
        v.literal("video"),
        v.literal("image"),
        v.literal("lyrics"),
        v.literal("webview"),
        v.literal("midi")
      )
    ),
    onlyLiked: v.optional(v.boolean()),
  },
  handler: async (ctx, args) => {
    const userId = await getAuthUserId(ctx);
    if (!userId) {
      return [];
    }

    // Get all atoms by this user
    let atoms = await ctx.db
      .query("atoms")
      .withIndex("by_owner", (q) => q.eq("ownerId", userId))
      .order("desc")
      .collect();

    // Filter by type if specified
    if (args.type !== undefined) {
      atoms = atoms.filter((atom) => atom.type === args.type);
    }

    // Filter by liked status if specified
    if (args.onlyLiked) {
      atoms = atoms.filter((atom) => atom.likedBy?.includes(userId));
    }

    return atoms;
  },
});

// Generate upload URL for user images
export const generateImageUploadUrl = mutation(async (ctx) => {
  const userId = await getAuthUserId(ctx);
  if (!userId) throw new Error("Not authenticated");
  return await ctx.storage.generateUploadUrl();
});

// Create image atom from uploaded image
export const createImageAtom = mutation({
  args: {
    storageId: v.id("_storage"),
    spaceId: v.id("spaces"),
    width: v.optional(v.number()),
    height: v.optional(v.number()),
    filename: v.optional(v.string()),
  },
  handler: async (ctx, args) => {
    const userId = await getAuthUserId(ctx);
    if (!userId) throw new Error("Not authenticated");

    // Verify user is a member of the space
    const member = await ctx.db
      .query("spaceMembers")
      .withIndex("by_space_and_user", (q) =>
        q.eq("spaceId", args.spaceId).eq("userId", userId)
      )
      .first();

    if (!member || member.isKicked) {
      throw new Error("Not authorized - not a member of this space");
    }

    const atomId = await ctx.db.insert("atoms", {
      type: "image",
      spaceId: args.spaceId,
      ownerId: userId,
      metadata: {
        storageId: args.storageId,
        width: args.width,
        height: args.height,
        filename: args.filename,
      },
      status: "completed",
      progress: 100,
      createdAt: Date.now(),
      updatedAt: Date.now(),
    });

    return atomId;
  },
});

// ============================================================================
// PROVENANCE SYSTEM
// ============================================================================

// Create an operation record
export const createOperation = internalMutation({
  args: {
    type: v.union(
      v.literal("generate"),
      v.literal("cover"),
      v.literal("stem"),
      v.literal("midi"),
      v.literal("remix"),
      v.literal("mashup"),
      v.literal("style_transfer"),
      v.literal("upload"),
      v.literal("continue")
    ),
    parameters: v.optional(v.any()),
    createdBy: v.id("users"),
  },
  handler: async (ctx, args) => {
    return await ctx.db.insert("operations", {
      type: args.type,
      parameters: args.parameters,
      createdAt: Date.now(),
      createdBy: args.createdBy,
    });
  },
});

// Link an atom to an operation (as input or output)
export const createOperationAtom = internalMutation({
  args: {
    operationId: v.id("operations"),
    atomId: v.id("atoms"),
    direction: v.union(v.literal("input"), v.literal("output")),
    role: v.optional(v.string()),
  },
  handler: async (ctx, args) => {
    return await ctx.db.insert("operationAtoms", {
      operationId: args.operationId,
      atomId: args.atomId,
      direction: args.direction,
      role: args.role,
    });
  },
});

// Get provenance for an atom (how it was created)
export const getProvenanceByAtom = query({
  args: {
    atomId: v.id("atoms"),
  },
  handler: async (ctx, args) => {
    const userId = await getAuthUserId(ctx);
    if (!userId) throw new Error("Not authenticated");

    // Find the output record for this atom
    const outputRecord = await ctx.db
      .query("operationAtoms")
      .withIndex("by_atom_and_direction", (q) =>
        q.eq("atomId", args.atomId).eq("direction", "output")
      )
      .first();

    if (!outputRecord) return null;

    // Get the operation
    const operation = await ctx.db.get(outputRecord.operationId);
    if (!operation) return null;

    // Get all inputs for this operation
    const inputRecords = await ctx.db
      .query("operationAtoms")
      .withIndex("by_operation_and_direction", (q) =>
        q.eq("operationId", outputRecord.operationId).eq("direction", "input")
      )
      .collect();

    // Load the actual input atoms
    const inputs = await Promise.all(
      inputRecords.map(async (input) => {
        const atom = await ctx.db.get(input.atomId);
        return {
          atomId: input.atomId,
          role: input.role,
          atom,
        };
      })
    );

    return {
      operation,
      inputs,
      outputRole: outputRecord.role,
    };
  },
});

// Get all children/derivatives of an atom
export const getAtomChildren = query({
  args: {
    atomId: v.id("atoms"),
  },
  handler: async (ctx, args) => {
    const userId = await getAuthUserId(ctx);
    if (!userId) throw new Error("Not authenticated");

    // Find all operations where this atom was an input
    const inputRecords = await ctx.db
      .query("operationAtoms")
      .withIndex("by_atom_and_direction", (q) =>
        q.eq("atomId", args.atomId).eq("direction", "input")
      )
      .collect();

    // For each operation, get its outputs
    const children = await Promise.all(
      inputRecords.map(async (inputRec) => {
        const operation = await ctx.db.get(inputRec.operationId);
        if (!operation) return null;

        // Get output atoms for this operation
        const outputs = await ctx.db
          .query("operationAtoms")
          .withIndex("by_operation_and_direction", (q) =>
            q.eq("operationId", inputRec.operationId).eq("direction", "output")
          )
          .collect();

        // Load actual output atoms
        const outputAtoms = await Promise.all(
          outputs.map(async (out) => {
            const atom = await ctx.db.get(out.atomId);
            return { atom, outputRole: out.role };
          })
        );

        return {
          operation,
          inputRole: inputRec.role,
          outputs: outputAtoms.filter((o) => o.atom !== null),
        };
      })
    );

    return children.filter((c) => c !== null);
  },
});

// ============================================================================
// STEM GENERATION
// ============================================================================

// Generate stems for a source song
export const generateStems = internalAction({
  args: {
    sourceAtomId: v.id("atoms"),
    requestedStems: v.optional(v.array(v.string())),
    responseMessageId: v.optional(v.id("messages")),
    spaceId: v.id("spaces"),
    userId: v.id("users"),
  },
  handler: async (ctx, args): Promise<Id<"atoms">[]> => {
    // Get user's Suno token if they have one
    const userToken = await ctx.runQuery(internal.users.getUserTokenInternal, {
      userId: args.userId,
    });

    const client = createSunoClient(userToken);

    try {
      // 1. Get source atom and extract Suno clip ID
      const sourceAtom: any = await ctx.runQuery(internal.atoms.getAtomInternal, {
        id: args.sourceAtomId,
      });

      if (!sourceAtom || sourceAtom.type !== "song") {
        throw new Error("Source atom must be a song");
      }

      const sunoClipId: string = sourceAtom.metadata?.sunoClipId;
      if (!sunoClipId) {
        throw new Error("Source song has no Suno clip ID");
      }

      // 2. Generate stems via Suno API (returns 24 stems - 2 sets of 12)
      // We only need the first 12 stems (first half of response)
      const allStemClipIds = await client.generateStems(sunoClipId);
      const stemClipIds = allStemClipIds.slice(0, Math.floor(allStemClipIds.length / 2));

      // 3. Create operation record
      const operationId = await ctx.runMutation(internal.atoms.createOperation, {
        type: "stem",
        parameters: {
          stemTask: "twelve",
          stemTypeGroupName: "Twelve",
          requestedStems: args.requestedStems,
        },
        createdBy: args.userId,
      });

      // 4. Create input link (source song)
      await ctx.runMutation(internal.atoms.createOperationAtom, {
        operationId,
        atomId: args.sourceAtomId,
        direction: "input",
        role: "source",
      });

      // 5. Create stem atoms for each clip (all 12 initially, will filter later)
      const stemAtomIds: Id<"atoms">[] = await Promise.all(
        stemClipIds.map(async (clipId, index): Promise<Id<"atoms">> =>
          ctx.runMutation(internal.atoms.createInternal, {
            type: "song",
            spaceId: args.spaceId,
            ownerId: args.userId,
            metadata: {
              sunoClipId: clipId,
              isStem: true,
              stemType: `stem_${index}`, // Will be updated when polling
              title: `${sourceAtom.metadata?.title || "Song"} - Stem ${index + 1}`,
            },
            status: "pending",
            progress: 0,
          })
        )
      );

      // 6. Create output links for all stems
      await Promise.all(
        stemAtomIds.map((stemAtomId: Id<"atoms">) =>
          ctx.runMutation(internal.atoms.createOperationAtom, {
            operationId,
            atomId: stemAtomId,
            direction: "output",
            role: "primary",
          })
        )
      );

      // 7. Schedule polling with filtering logic
      await ctx.scheduler.runAfter(0, internal.atoms.pollStemStatus, {
        atomIds: stemAtomIds,
        requestedStems: args.requestedStems,
        responseMessageId: args.responseMessageId,
        sourceAtomId: args.sourceAtomId,
      });

      return stemAtomIds;
    } catch (error) {
      console.error("Stem generation failed:", error);
      throw error;
    }
  },
});

// Poll stem status
export const pollStemStatus = internalAction({
  args: {
    atomIds: v.array(v.id("atoms")),
    requestedStems: v.optional(v.array(v.string())),
    responseMessageId: v.optional(v.id("messages")),
    sourceAtomId: v.optional(v.id("atoms")),
  },
  handler: async (ctx, args) => {
    // Get atoms
    const atoms = await Promise.all(
      args.atomIds.map((id) =>
        ctx.runQuery(internal.atoms.getAtomInternal, { id })
      )
    );

    // Filter active atoms
    const activeAtoms = atoms.filter(
      (atom: any) => atom && !["completed", "failed"].includes(atom.status)
    );

    if (activeAtoms.length === 0) {
      // All stems are complete or failed
      // If requestedStems was provided, filter the stems now
      if (args.requestedStems && args.requestedStems.length > 0) {
        await ctx.runAction(internal.atoms.filterStemsByRequest, {
          atomIds: args.atomIds,
          requestedStems: args.requestedStems,
          responseMessageId: args.responseMessageId,
        });
      } else if (args.responseMessageId) {
        // No filtering needed - update message with all stems
        const completedStems = atoms.filter(
          (atom: any) => atom && atom.status === "completed"
        );
        const stemIds = completedStems.map((atom: any) => atom._id);

        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎼 Here are your stems! They're ready.`,
          atomReferences: stemIds,
          assistantStatus: "completed",
        });
      }
      return;
    }

    // Get source atom to access lyrics (only fetch once)
    let sourceAtom: any = null;
    if (args.sourceAtomId) {
      sourceAtom = await ctx.runQuery(internal.atoms.getAtomInternal, {
        id: args.sourceAtomId,
      });
    }

    // Get user token from first atom's owner (all atoms in a batch share the same owner)
    const userId = activeAtoms[0]?.ownerId;
    const userToken = userId
      ? await ctx.runQuery(internal.users.getUserTokenInternal, { userId })
      : null;

    // Poll Suno API
    const client = createSunoClient(userToken);
    const sunoIds = activeAtoms.map((a: any) => a.metadata.sunoClipId);

    try {
      const stems = await client.getSongStatus(sunoIds);

      // Update atoms
      for (let i = 0; i < activeAtoms.length; i++) {
        const atom = activeAtoms[i];
        if (!atom) continue;
        const stem = stems[i];

        // Check if this is a vocals stem by looking at the title
        const isVocalsStem = stem.title?.toLowerCase().includes("vocal");

        // Prepare metadata update
        const updatedMetadata: any = {
          ...atom.metadata,
          title: stem.title || atom.metadata.title,
          artist: stem.artist,
          audioUrl: stem.audioUrl,
          videoUrl: stem.videoUrl,
          albumArtUrl: stem.albumArtUrl,
          duration: stem.duration,
          sunoStatus: stem.status,
        };

        // If this is a vocals stem and we have source lyrics, copy them
        // Check both metadata.lyrics and metadata.prompt for backwards compatibility
        if (isVocalsStem && sourceAtom?.metadata) {
          const sourceLyrics = sourceAtom.metadata.lyrics || sourceAtom.metadata.prompt;
          if (sourceLyrics) {
            updatedMetadata.lyrics = sourceLyrics;
          }
        }

        await ctx.runMutation(internal.atoms.updateInternal, {
          id: atom._id,
          metadata: updatedMetadata,
          status: stem.status,
          progress:
            stem.status === "completed"
              ? 100
              : stem.status === "streaming"
                ? 75
                : stem.status === "processing"
                  ? 50
                  : 25,
        });
      }

      // Schedule next poll if needed
      const stillActive = stems.some(
        (s) => !["completed", "failed"].includes(s.status)
      );

      if (stillActive) {
        await ctx.scheduler.runAfter(2000, internal.atoms.pollStemStatus, {
          atomIds: args.atomIds,
          requestedStems: args.requestedStems,
          responseMessageId: args.responseMessageId,
          sourceAtomId: args.sourceAtomId,
        });
      } else {
        // All stems complete - filter if needed
        if (args.requestedStems && args.requestedStems.length > 0) {
          await ctx.runAction(internal.atoms.filterStemsByRequest, {
            atomIds: args.atomIds,
            requestedStems: args.requestedStems,
            responseMessageId: args.responseMessageId,
          });
        } else if (args.responseMessageId) {
          // No filtering needed - update message with all stems
          const completedStems = atoms.filter(
            (atom: any) => atom && atom.status === "completed"
          );
          const stemIds = completedStems.map((atom: any) => atom._id);

          await ctx.runMutation(internal.assistant.updateAssistantMessage, {
            messageId: args.responseMessageId,
            content: `🎼 Here are your stems! They're ready.`,
            atomReferences: stemIds,
            assistantStatus: "completed",
          });
        }
      }
    } catch (error) {
      console.error("Stem polling error:", error);
      // Retry polling after delay
      await ctx.scheduler.runAfter(2000, internal.atoms.pollStemStatus, {
        atomIds: args.atomIds,
        requestedStems: args.requestedStems,
        responseMessageId: args.responseMessageId,
        sourceAtomId: args.sourceAtomId,
      });
    }
  },
});

// Filter stems based on user's request using OpenAI
export const filterStemsByRequest = internalAction({
  args: {
    atomIds: v.array(v.id("atoms")),
    requestedStems: v.array(v.string()),
    responseMessageId: v.optional(v.id("messages")),
  },
  handler: async (ctx, args) => {
    const openai = new OpenAI({
      apiKey: process.env.OPENAI_KEY,
    });

    // 1. Get all stem atoms with their titles
    const atoms = await Promise.all(
      args.atomIds.map((id) =>
        ctx.runQuery(internal.atoms.getAtomInternal, { id })
      )
    );

    // Filter to only completed atoms
    const completedAtoms = atoms.filter(
      (atom): atom is NonNullable<typeof atom> => atom !== null && atom !== undefined && atom.status === "completed"
    );

    if (completedAtoms.length === 0) {
      console.log("No completed stems to filter");
      return;
    }

    try {

      // 2. Build a list of stem titles for OpenAI to analyze
      const stemList = completedAtoms.map((atom: any, index: number) => ({
        index,
        atomId: atom._id,
        title: atom.metadata?.title || `Stem ${index + 1}`,
      }));

      const stemTitles = stemList.map((s) => `${s.index}: ${s.title}`).join("\n");

      // 3. Ask OpenAI to identify which stems match the user's request
      const prompt = `You are helping to filter music stems based on a user's request.

The user requested these stem types: ${args.requestedStems.join(", ")}

Here are the available stems:
${stemTitles}

Please identify which stems match the user's request. Look at the stem titles and determine which ones correspond to the requested types.

Respond with ONLY a JSON array of the indices (numbers) of the matching stems. For example: [0, 2, 5]

If a stem title contains words like "vocals", "voice", "singing", it matches "vocals".
If a stem title contains "drums", "percussion", "kick", "snare", it matches "drums".
If a stem title contains "bass", it matches "bass".
And so on for other common instrument types.

Be inclusive - if there's any reasonable match, include it.`;

      const response = await openai.chat.completions.create({
        model: "gpt-4o-mini",
        messages: [
          {
            role: "user",
            content: prompt,
          },
        ],
        temperature: 0,
      });

      const content = response.choices[0].message.content || "[]";

      // Parse the response to get matching indices
      let matchingIndices: number[];
      try {
        matchingIndices = JSON.parse(content);
      } catch (parseError) {
        console.error("Failed to parse OpenAI response:", content);
        // Fallback: keep all stems
        return;
      }

      // 4. Delete non-matching stems
      const matchingAtomIds = new Set(
        matchingIndices.map((i) => stemList[i]?.atomId).filter(Boolean)
      );

      for (const atom of completedAtoms) {
        if (!matchingAtomIds.has(atom._id)) {
          console.log(`Deleting non-matching stem: ${atom.metadata?.title}`);
          await ctx.runMutation(internal.atoms.deleteAtom, {
            atomId: atom._id,
          });
        }
      }

      console.log(`Filtered stems: kept ${matchingAtomIds.size} out of ${completedAtoms.length}`);

      // Update the assistant message with only the filtered stems
      if (args.responseMessageId) {
        const filteredStemIds = Array.from(matchingAtomIds);
        const stemsDescription = args.requestedStems.join(", ");

        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎼 Here are your ${stemsDescription} stem(s)! They're ready.`,
          atomReferences: filteredStemIds,
          assistantStatus: "completed",
        });
      }
    } catch (error) {
      console.error("Failed to filter stems:", error);
      // Don't throw - if filtering fails, keep all stems and update message anyway
      if (args.responseMessageId) {
        const allCompletedIds = args.atomIds.filter((id) =>
          completedAtoms.some((atom: any) => atom._id === id)
        );

        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎼 Here are your stems! (Note: filtering failed, showing all stems)`,
          atomReferences: allCompletedIds,
          assistantStatus: "completed",
        });
      }
    }
  },
});

// Delete an atom (internal mutation)
export const deleteAtom = internalMutation({
  args: {
    atomId: v.id("atoms"),
  },
  handler: async (ctx, args) => {
    await ctx.db.delete(args.atomId);
  },
});

// ============================================================================
// MIDI GENERATION
// ============================================================================

// Generate MIDI for a song or stem
export const generateMidi = internalAction({
  args: {
    sourceAtomId: v.id("atoms"),
    spaceId: v.id("spaces"),
    userId: v.id("users"),
  },
  handler: async (ctx, args): Promise<Id<"atoms">> => {
    try {
      // 1. Get source atom and extract Suno clip ID
      const sourceAtom: any = await ctx.runQuery(internal.atoms.getAtomInternal, {
        id: args.sourceAtomId,
      });

      if (!sourceAtom || sourceAtom.type !== "song") {
        throw new Error("Source atom must be a song");
      }

      const sunoClipId: string = sourceAtom.metadata?.sunoClipId;
      if (!sunoClipId) {
        throw new Error("Source song has no Suno clip ID");
      }

      // 2. Create operation record
      const operationId = await ctx.runMutation(internal.atoms.createOperation, {
        type: "midi",
        parameters: {},
        createdBy: args.userId,
      });

      // 3. Create input link
      await ctx.runMutation(internal.atoms.createOperationAtom, {
        operationId,
        atomId: args.sourceAtomId,
        direction: "input",
        role: "source",
      });

      // 4. Create MIDI atom
      const midiAtomId: Id<"atoms"> = await ctx.runMutation(internal.atoms.createInternal, {
        type: "midi",
        spaceId: args.spaceId,
        ownerId: args.userId,
        metadata: {
          sunoClipId,
          title: `${sourceAtom.metadata?.title || "Song"} - MIDI`,
        },
        status: "pending",
        progress: 0,
      });

      // 5. Create output link
      await ctx.runMutation(internal.atoms.createOperationAtom, {
        operationId,
        atomId: midiAtomId,
        direction: "output",
        role: "primary",
      });

      // 6. Schedule polling (start immediately)
      await ctx.scheduler.runAfter(0, internal.atoms.pollMidiStatus, {
        atomId: midiAtomId,
        retryCount: 0,
      });

      return midiAtomId;
    } catch (error) {
      console.error("MIDI generation failed:", error);
      throw error;
    }
  },
});

// Poll MIDI generation status
export const pollMidiStatus = internalAction({
  args: {
    atomId: v.id("atoms"),
    retryCount: v.optional(v.number()),
  },
  handler: async (ctx, args) => {
    const MAX_RETRIES = 36; // 36 retries * 5s = 180s (3 minutes) max
    const currentRetry = args.retryCount || 0;

    // Safety check: stop if we've exceeded max retries
    if (currentRetry >= MAX_RETRIES) {
      console.error(`MIDI polling exceeded max retries (${MAX_RETRIES}) for atom ${args.atomId}`);
      await ctx.runMutation(internal.atoms.updateInternal, {
        id: args.atomId,
        status: "failed",
        metadata: { error: "Polling timeout - exceeded max retries" },
      });
      return;
    }

    // Get MIDI atom
    const atom = await ctx.runQuery(internal.atoms.getAtomInternal, {
      id: args.atomId,
    });

    if (!atom || ["completed", "failed"].includes(atom.status)) {
      console.log(`MIDI atom ${args.atomId} is already in terminal state: ${atom?.status}`);
      return; // Already done
    }

    const sunoClipId = atom.metadata?.sunoClipId;
    if (!sunoClipId) {
      console.error(`MIDI atom ${args.atomId} has no Suno clip ID`);
      await ctx.runMutation(internal.atoms.updateInternal, {
        id: args.atomId,
        status: "failed",
        metadata: { error: "No Suno clip ID" },
      });
      return;
    }

    // Get user token from atom owner
    const userId = atom.ownerId;
    const userToken = userId
      ? await ctx.runQuery(internal.users.getUserTokenInternal, { userId })
      : null;

    const client = createSunoClient(userToken);

    try {
      const midiStatus = await client.getMidiStatus(sunoClipId);
      console.log("MIDI status response for", sunoClipId, ":", JSON.stringify(midiStatus, null, 2));

      // Update MIDI atom
      await ctx.runMutation(internal.atoms.updateInternal, {
        id: args.atomId,
        metadata: {
          ...atom.metadata,
          midiUrl: midiStatus.midiUrl,
          midiData: midiStatus.midiData,
          midiStatus: midiStatus.status,
        },
        status: midiStatus.status,
        progress:
          midiStatus.status === "completed"
            ? 100
            : midiStatus.status === "processing"
              ? 50
              : 25,
      });

      // Schedule next poll if not complete (5s interval for MIDI)
      if (!["completed", "failed"].includes(midiStatus.status)) {
        await ctx.scheduler.runAfter(5000, internal.atoms.pollMidiStatus, {
          atomId: args.atomId,
          retryCount: currentRetry + 1,
        });
      }
    } catch (error) {
      console.error("MIDI polling error:", error);

      // Only retry if under max retries
      if (currentRetry < MAX_RETRIES - 1) {
        await ctx.scheduler.runAfter(5000, internal.atoms.pollMidiStatus, {
          atomId: args.atomId,
          retryCount: currentRetry + 1,
        });
      } else {
        // Mark as failed if we've hit max retries
        await ctx.runMutation(internal.atoms.updateInternal, {
          id: args.atomId,
          status: "failed",
          metadata: { error: "Polling failed after max retries" },
        });
      }
    }
  },
});