import { v } from "convex/values";
import OpenAI from "openai";
import { internal } from "./_generated/api";
import { Id } from "./_generated/dataModel";
import { internalAction, internalMutation } from "./_generated/server";

// Assistant mention constant
export const ASSISTANT_MENTION = "suno";

// System prompt for Orphy
const ASSISTANT_SYSTEM_PROMPT = `You are Orphy, an AI assistant for Suno Spaces. Your job is to help users create music and visual content.

When a user mentions @orphy for song generation, you must:

1. **Generate lyrics** (for the 'lyrics' field):
   - If user provides explicit lyrics, use them as-is
   - If user describes a theme/topic (e.g., "summer vibes", "heartbreak"), write original lyrics about that theme
   - If user wants instrumental, lyrics can be a short instrumental description

   **IMPORTANT FORMATTING RULES**:
   - Use section tags: [Verse], [Chorus], [PreChorus], [Bridge], [Outro]
   - Structure: Verse → PreChorus → Chorus, then repeat variations
   - Make choruses catchy, memorable, and hook-driven
   - Avoid excessive rhyming - natural flow is better than forced rhymes
   - Use conversational language, not overly poetic

   **Example structure**:
   [Verse]
   Walking down the street at midnight
   City lights reflecting in my eyes

   [PreChorus]
   Something's changing, I can feel it

   [Chorus]
   This is where we come alive
   Dancing under neon skies

   [Verse 2]
   Empty cafes, quiet conversations
   ...

2. **Extract genre/style tags** (for the 'tags' field):
   - Parse genre/style descriptors from user's message
   - Format as comma-separated string (e.g., "indie rock, energetic, 2000s")
   - Infer appropriate styles if not explicitly stated

3. **Determine if instrumental** (for the 'makeInstrumental' field):
   - Set to true only if user explicitly requests instrumental

4. **Detect cover/remix requests** (for the 'coverAtomId' field):
   - If user references a song atom and asks to change style/genre (e.g., "make this jazz", "turn this into metal")
   - Look for atom IDs in the message context (they will be provided in brackets like "[Referenced atom IDs: jd7abc123...]")
   - Extract the atom ID and provide it in 'coverAtomId'
   - The 'tags' field should contain the NEW style/genre for the cover
   - Atom IDs are alphanumeric strings (e.g., "jd7abc123def456")

Examples:
- "create a chill lofi beat" → prompt: "[instrumental description]", tags: "lofi, chill, instrumental", makeInstrumental: true
- "make a song about summer" → prompt: "[Verse]\n...\n[PreChorus]\n...\n[Chorus]\n...", tags: "pop, summer, upbeat"
- "indie rock song about leaving home" → prompt: "[Verse]\n...\n[Chorus]\n...", tags: "indie rock, emotional, storytelling"
- "make this jazz [Referenced atom IDs: jd7abc123def456]" → prompt: "", tags: "jazz, smooth", coverAtomId: "jd7abc123def456"

For image generation requests, use the generate_image tool:
- Extract or create a detailed image description from the user's request
- Be specific about style, composition, colors, and mood
- Examples: "generate an album cover with neon lights", "create an image of a sunset over mountains"`;

// Tool definition for OpenAI function calling
const GENERATE_SONG_TOOL: OpenAI.Chat.Completions.ChatCompletionTool = {
  type: "function",
  function: {
    name: "generate_song",
    description:
      "Generate a new song. You must write original lyrics or an instrumental description. Can also create a cover/remix of an existing song by changing its style/genre.",
    parameters: {
      type: "object",
      properties: {
        lyrics: {
          type: "string",
          description:
            "The actual lyrics to be sung, or instrumental description. CRITICAL FORMATTING: Use section tags like [Verse], [Chorus], [PreChorus], [Bridge]. Structure should be Verse → PreChorus → Chorus. Make choruses catchy and memorable. Avoid excessive rhyming - natural flow is better. Example format:\n[Verse]\nWalking down the street at midnight\nCity lights reflecting in my eyes\n\n[PreChorus]\nSomething's changing, I can feel it\n\n[Chorus]\nThis is where we come alive\nDancing under neon skies",
        },
        tags: {
          type: "string",
          description:
            'Comma-separated genre/style tags for musical style (e.g., "indie rock, energetic, 2000s", "lofi, chill, jazz")',
        },
        makeInstrumental: {
          type: "boolean",
          description: "Whether to create an instrumental version (no vocals)",
        },
        coverAtomId: {
          type: "string",
          description: "Optional: The atom ID of an existing song to create a cover/remix of. When provided, the song will be recreated in the style specified by 'tags'. Look for atom IDs in the format '[Referenced atom IDs: ...]' in the user message. Example: 'jd7abc123def456' (alphanumeric string, NOT 'suno' or 'orphy').",
        },
      },
      required: ["lyrics", "tags"],
    },
  },
};

const GENERATE_IMAGE_TOOL: OpenAI.Chat.Completions.ChatCompletionTool = {
  type: "function",
  function: {
    name: "generate_image",
    description:
      "Generate images using AI. Creates 2 variations of the image based on the prompt.",
    parameters: {
      type: "object",
      properties: {
        prompt: {
          type: "string",
          description:
            "Detailed description of the image to generate. Be specific about style, composition, colors, mood, and subject matter.",
        },
      },
      required: ["prompt"],
    },
  },
};

const GENERATE_STEMS_TOOL: OpenAI.Chat.Completions.ChatCompletionTool = {
  type: "function",
  function: {
    name: "generate_stems",
    description:
      "Generate stem tracks (isolated instruments/vocals) from a song. Can generate all stems or specific stems based on user request. Requires a source song atom ID.",
    parameters: {
      type: "object",
      properties: {
        sourceAtomId: {
          type: "string",
          description:
            "The atom ID of the source song to generate stems from. Look for atom IDs in the format '[Referenced atom IDs: ...]' in the user message. Example: 'jd7abc123def456' (alphanumeric string, NOT 'suno' or 'orphy').",
        },
        requestedStems: {
          type: "array",
          items: { type: "string" },
          description:
            "Optional: Array of specific stem types requested by the user (e.g., ['vocals', 'drums', 'bass']). If user asks for 'all stems' or doesn't specify, omit this field to return all stems. Common stem types: vocals, drums, bass, guitar, piano, keys, strings, synth, instrumental, other.",
        },
      },
      required: ["sourceAtomId"],
    },
  },
};

const GENERATE_MIDI_TOOL: OpenAI.Chat.Completions.ChatCompletionTool = {
  type: "function",
  function: {
    name: "generate_midi",
    description:
      "Generate MIDI data from a song or stem. Requires a source atom ID.",
    parameters: {
      type: "object",
      properties: {
        sourceAtomId: {
          type: "string",
          description:
            "The atom ID of the source song or stem to generate MIDI from. Look for atom IDs in the format '[Referenced atom IDs: ...]' in the user message. Example: 'jd7abc123def456' (alphanumeric string, NOT 'suno' or 'orphy').",
        },
      },
      required: ["sourceAtomId"],
    },
  },
};

// Handle incoming message with @suno mention
export const handleMessage = internalMutation({
  args: {
    messageId: v.id("messages"),
  },
  handler: async (ctx, args) => {
    const message = await ctx.db.get(args.messageId);
    if (!message) throw new Error("Message not found");

    console.log("handling message", message);

    // Check if message mentions @suno (or legacy @orphy)
    const mentionsAssistant = message.mentions?.includes(ASSISTANT_MENTION as any) ||
                              message.mentions?.includes("orphy" as any);
    if (!mentionsAssistant) return;

    // Get room and space context
    const room = await ctx.db.get(message.roomId);
    if (!room) throw new Error("Room not found");

    // Create immediate "thinking" response message
    const responseMessageId = await ctx.runMutation(internal.assistant.createAssistantMessage, {
      roomId: message.roomId,
      content: "🤔 Thinking...",
      assistantStatus: "processing",
      replyToId: message._id,
    });

    // Schedule assistant action to process the request
    await ctx.scheduler.runAfter(
      0,
      internal.assistant.processAssistantRequest,
      {
        messageId: message._id,
        responseMessageId,
        roomId: message.roomId,
        spaceId: room.spaceId,
        userId: message.userId,
        content: message.content,
        atomReferences: message.atomReferences,
      }
    );
  },
});

// Process assistant request (internal action)
export const processAssistantRequest = internalAction({
  args: {
    messageId: v.id("messages"),
    responseMessageId: v.id("messages"),
    roomId: v.id("rooms"),
    spaceId: v.id("spaces"),
    userId: v.id("users"),
    content: v.string(),
    atomReferences: v.optional(v.array(v.id("atoms"))),
  },
  handler: async (ctx, args) => {
    try {
      // 1. Classify intent using OpenAI
      const result = await ctx.runAction(internal.assistant.classifyIntent, {
        messageContent: args.content,
        atomReferences: args.atomReferences,
      });

      if (!result) {
        // Not a recognized request - update the thinking message
        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content:
            "I can help you create songs and images! Try: @suno create a chill lofi beat, or @suno generate an album cover",
          assistantStatus: "completed",
        });
        return;
      }

      // 2. Handle based on tool type
      if (result.toolName === "generate_song") {
        const params = result.params as {
          lyrics: string;
          tags?: string;
          makeInstrumental?: boolean;
          coverAtomId?: string;
        };

        // Get cover clip ID and original lyrics if coverAtomId is provided
        let coverClipId: string | undefined;
        let originalPrompt = params.lyrics;
        if (params.coverAtomId) {
          // Validate that coverAtomId looks like a valid Convex ID
          // Convex IDs start with the table name prefix (e.g., "jd7..." for atoms)
          const isValidId = params.coverAtomId.match(/^[a-z0-9]+$/i) &&
                           !params.coverAtomId.includes("suno") &&
                           !params.coverAtomId.includes("orphy");

          if (isValidId) {
            try {
              const coverAtom = await ctx.runQuery(internal.atoms.getAtomInternal, {
                id: params.coverAtomId as Id<"atoms">,
              });
              if (coverAtom && coverAtom.type === "song" && coverAtom.metadata?.sunoClipId) {
                coverClipId = coverAtom.metadata.sunoClipId;

                // TODO: Eventually, allow OpenAI to modify lyrics for covers.
                // For now, we override and use the original song's lyrics/prompt.
                if (coverAtom.metadata?.prompt) {
                  originalPrompt = coverAtom.metadata.prompt;
                }
              }
            } catch (error) {
              console.error("Failed to fetch cover atom:", error);
              // Continue without cover
            }
          } else {
            console.warn("Invalid coverAtomId provided by OpenAI:", params.coverAtomId);
          }
        }

        // Update message to show what we're doing
        const coverMessage = coverClipId ? " (cover version)" : "";
        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎵 Creating 2 songs${coverMessage}: "${params.tags}". I'll let you know when they're ready!`,
          assistantStatus: "processing",
        });

        // Generate songs - pass lyrics as prompt to Suno API
        const atomIds = await ctx.runAction(internal.atoms.generateSong, {
          lyrics: originalPrompt,
          tags: params.tags,
          makeInstrumental: params.makeInstrumental,
          coverClipId,
          spaceId: args.spaceId,
          userId: args.userId,
        });

        // Attach atoms to response message with updated content
        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎵 Your songs are generating! They'll be ready soon.`,
          atomReferences: atomIds,
          assistantStatus: "completed",
        });

        // Schedule monitoring of song generation status to update the message
        await ctx.scheduler.runAfter(
          0,
          internal.assistant.monitorSongGenerationMessage,
          {
            messageId: args.responseMessageId,
            atomIds,
          }
        );
      } else if (result.toolName === "generate_image") {
        const params = result.params as { prompt: string };

        // Update message to show what we're doing
        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎨 Generating 2 images: "${params.prompt}". I'll let you know when they're ready!`,
          assistantStatus: "processing",
        });

        // Generate images
        const atomIds = await ctx.runAction(internal.atoms.generateImage, {
          prompt: params.prompt,
          spaceId: args.spaceId,
          userId: args.userId,
        });

        // Attach atoms to response message with updated content
        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎨 Your images are generating! They'll be ready soon.`,
          atomReferences: atomIds,
          assistantStatus: "completed",
        });
      } else if (result.toolName === "generate_stems") {
        const params = result.params as { sourceAtomId: string; requestedStems?: string[] };

        // Validate atom reference
        if (!args.atomReferences || args.atomReferences.length === 0) {
          await ctx.runMutation(internal.assistant.updateAssistantMessage, {
            messageId: args.responseMessageId,
            content: "⚠️ Please attach a song to generate stems.",
            assistantStatus: "failed",
          });
          return;
        }

        // Update message to show what we're doing
        const stemsDescription = params.requestedStems && params.requestedStems.length > 0
          ? params.requestedStems.join(", ")
          : "all";
        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎼 Generating ${stemsDescription} stem(s) for this song. I'll let you know when they're ready!`,
          assistantStatus: "processing",
        });

        // Generate stems (returns stem atom IDs for monitoring)
        const stemAtomIds = await ctx.runAction(internal.atoms.generateStems, {
          sourceAtomId: params.sourceAtomId as any,
          requestedStems: params.requestedStems,
          responseMessageId: args.responseMessageId,
          spaceId: args.spaceId,
          userId: args.userId,
        });

        // Attach initial stem atoms to response message
        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎼 Your stems are generating! They'll be ready soon.`,
          atomReferences: stemAtomIds,
          assistantStatus: "completed",
        });

        // Schedule monitoring of stem generation status
        // Note: The filtering and final message update will happen in the polling action
        await ctx.scheduler.runAfter(
          0,
          internal.assistant.monitorStemGenerationMessage,
          {
            messageId: args.responseMessageId,
            atomIds: stemAtomIds,
          }
        );
      } else if (result.toolName === "generate_midi") {
        const params = result.params as { sourceAtomId: string };

        // Validate atom reference
        if (!args.atomReferences || args.atomReferences.length === 0) {
          await ctx.runMutation(internal.assistant.updateAssistantMessage, {
            messageId: args.responseMessageId,
            content: "⚠️ Please attach a song to generate MIDI.",
            assistantStatus: "failed",
          });
          return;
        }

        // Update message to show what we're doing
        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎹 Generating MIDI. I'll let you know when it's ready!`,
          assistantStatus: "processing",
        });

        // Generate MIDI
        const midiAtomId = await ctx.runAction(internal.atoms.generateMidi, {
          sourceAtomId: params.sourceAtomId as any,
          spaceId: args.spaceId,
          userId: args.userId,
        });

        // Attach MIDI to response message with updated content
        await ctx.runMutation(internal.assistant.updateAssistantMessage, {
          messageId: args.responseMessageId,
          content: `🎹 Here's your MIDI file! It'll finish processing shortly.`,
          atomReferences: [midiAtomId],
          assistantStatus: "completed",
        });
      }
    } catch (error) {
      console.error("Assistant request failed:", error);

      // Update the existing thinking message with error
      await ctx.runMutation(internal.assistant.updateAssistantMessage, {
        messageId: args.responseMessageId,
        content: "😔 Something went wrong. Let's try again?",
        assistantStatus: "failed",
      });
    }
  },
});

// Classify intent using OpenAI (internal action)
export const classifyIntent = internalAction({
  args: {
    messageContent: v.string(),
    atomReferences: v.optional(v.array(v.id("atoms"))),
  },
  handler: async (ctx, args) => {
    const openai = new OpenAI({
      apiKey: process.env.OPENAI_KEY,
    });

    try {
      // Build user message with atom references if present
      let userMessage = args.messageContent;
      if (args.atomReferences && args.atomReferences.length > 0) {
        const atomIds = args.atomReferences.join(", ");
        userMessage += `\n\n[Referenced atom IDs: ${atomIds}]`;
      }

      const response = await openai.chat.completions.create({
        model: "gpt-4o-mini",
        messages: [
          {
            role: "system",
            content: ASSISTANT_SYSTEM_PROMPT,
          },
          {
            role: "user",
            content: userMessage,
          },
        ],
        tools: [
          GENERATE_SONG_TOOL,
          GENERATE_IMAGE_TOOL,
          GENERATE_STEMS_TOOL,
          GENERATE_MIDI_TOOL,
        ],
        tool_choice: "auto",
      });

      const toolCall = response.choices[0].message.tool_calls?.[0];
      if (!toolCall) {
        return null;
      }

      return {
        toolName: toolCall.function.name,
        params: JSON.parse(toolCall.function.arguments),
      };
    } catch (error) {
      console.error("OpenAI classification failed:", error);
      return null;
    }
  },
});

// Create assistant message (internal mutation)
export const createAssistantMessage = internalMutation({
  args: {
    roomId: v.id("rooms"),
    content: v.string(),
    assistantStatus: v.union(
      v.literal("processing"),
      v.literal("completed"),
      v.literal("failed")
    ),
    atomReferences: v.optional(v.array(v.id("atoms"))),
    replyToId: v.optional(v.id("messages")),
  },
  handler: async (ctx, args) => {
    // Create a system user ID for Orphy
    // For now, use the first user in the system, but this should be a dedicated system user
    const users = await ctx.db.query("users").first();
    const orphyUserId = users?._id;

    if (!orphyUserId) {
      throw new Error("No users in system to act as Orphy");
    }

    return await ctx.db.insert("messages", {
      roomId: args.roomId,
      userId: orphyUserId,
      content: args.content,
      timestamp: Date.now(),
      isAssistantMessage: true,
      assistantStatus: args.assistantStatus,
      atomReferences: args.atomReferences,
      replyToId: args.replyToId,
    });
  },
});

// Update assistant message (internal mutation)
export const updateAssistantMessage = internalMutation({
  args: {
    messageId: v.id("messages"),
    content: v.optional(v.string()),
    atomReferences: v.optional(v.array(v.id("atoms"))),
    assistantStatus: v.optional(
      v.union(
        v.literal("processing"),
        v.literal("completed"),
        v.literal("failed")
      )
    ),
  },
  handler: async (ctx, args) => {
    const { messageId, ...updates } = args;
    await ctx.db.patch(messageId, updates);
  },
});

// Monitor song generation and update message accordingly (internal action)
export const monitorSongGenerationMessage = internalAction({
  args: {
    messageId: v.id("messages"),
    atomIds: v.array(v.id("atoms")),
  },
  handler: async (ctx, args) => {
    // Get current atom statuses
    const atoms = await Promise.all(
      args.atomIds.map((id) =>
        ctx.runQuery(internal.atoms.getAtomInternal, { id })
      )
    );

    // Count statuses
    const statuses = {
      pending: 0,
      processing: 0,
      streaming: 0,
      completed: 0,
      failed: 0,
    };

    for (const atom of atoms) {
      if (atom) {
        statuses[atom.status as keyof typeof statuses]++;
      }
    }

    const total = atoms.length;
    const allDone = statuses.completed + statuses.failed === total;
    const anyStreaming = statuses.streaming > 0;
    const anyProcessing = statuses.processing > 0 || statuses.pending > 0;

    // Determine new message content based on statuses
    let newContent: string;
    if (allDone) {
      // All songs are done
      if (statuses.failed === total) {
        newContent = "😔 Something went wrong with song generation.";
      } else if (statuses.failed > 0) {
        newContent = `✅ ${statuses.completed} song(s) ready! ${statuses.failed} failed.`;
      } else {
        newContent = `✅ All ${total} songs are ready!`;
      }
    } else if (anyStreaming && !anyProcessing) {
      // All songs are either streaming or completed (ready to play but some still finalizing)
      newContent = `🎵 ${statuses.streaming + statuses.completed} song(s) ready to play! ${statuses.streaming > 0 ? "Still finalizing..." : ""}`;
    } else if (anyStreaming) {
      // Mix of streaming and processing
      const readyCount = statuses.streaming + statuses.completed;
      newContent = `🎵 ${readyCount} song(s) ready to play! ${total - readyCount} still generating...`;
    } else {
      // Still generating
      newContent = "🎵 Your songs are generating! They'll be ready soon.";
    }

    // Update the message
    await ctx.runMutation(internal.assistant.updateAssistantMessage, {
      messageId: args.messageId,
      content: newContent,
    });

    // Schedule next check if not all done
    if (!allDone) {
      await ctx.scheduler.runAfter(
        2000,
        internal.assistant.monitorSongGenerationMessage,
        {
          messageId: args.messageId,
          atomIds: args.atomIds,
        }
      );
    }
  },
});

// Monitor stem generation and update message accordingly (internal action)
export const monitorStemGenerationMessage = internalAction({
  args: {
    messageId: v.id("messages"),
    atomIds: v.array(v.id("atoms")),
  },
  handler: async (ctx, args) => {
    // Get current atom statuses
    const atoms = await Promise.all(
      args.atomIds.map((id) =>
        ctx.runQuery(internal.atoms.getAtomInternal, { id })
      )
    );

    // Count statuses
    const statuses = {
      pending: 0,
      processing: 0,
      streaming: 0,
      completed: 0,
      failed: 0,
    };

    for (const atom of atoms) {
      if (atom) {
        statuses[atom.status as keyof typeof statuses]++;
      }
    }

    const total = atoms.length;
    const allDone = statuses.completed + statuses.failed === total;
    const anyStreaming = statuses.streaming > 0;
    const anyProcessing = statuses.processing > 0 || statuses.pending > 0;

    // Determine new message content based on statuses
    let newContent: string;
    if (allDone) {
      if (statuses.failed === total) {
        newContent = "😔 Something went wrong with stem generation.";
      } else if (statuses.failed > 0) {
        newContent = `✅ ${statuses.completed} stem(s) ready! ${statuses.failed} failed.`;
      } else {
        newContent = `✅ All ${total} stems are ready!`;
      }
    } else if (anyStreaming && !anyProcessing) {
      newContent = `🎼 ${statuses.streaming + statuses.completed} stem(s) ready to play! ${statuses.streaming > 0 ? "Still finalizing..." : ""}`;
    } else if (anyStreaming) {
      const readyCount = statuses.streaming + statuses.completed;
      newContent = `🎼 ${readyCount} stem(s) ready to play! ${total - readyCount} still generating...`;
    } else {
      newContent = "🎼 Your stems are generating! They'll be ready soon.";
    }

    // Update the message
    await ctx.runMutation(internal.assistant.updateAssistantMessage, {
      messageId: args.messageId,
      content: newContent,
    });

    // Schedule next check if not all done
    if (!allDone) {
      await ctx.scheduler.runAfter(
        2000,
        internal.assistant.monitorStemGenerationMessage,
        {
          messageId: args.messageId,
          atomIds: args.atomIds,
        }
      );
    }
  },
});