import { createTool } from '@mastra/core/tools';
import { z } from 'zod';
import { createMastraCode } from 'mastracode';
import type { HarnessMessage } from '@mastra/core/harness';

const inputSchema = z.object({
  projectPath: z
    .string()
    .describe('Absolute path to the project root (e.g. from repositoryListTool directory)'),
  query: z
    .string()
    .describe('The user question or task to send to Mastra Code for analysis'),
  threadId: z
    .string()
    .optional()
    .describe(
      'Mastra Code thread ID to use. For the first call when you want a new thread, you may omit this field or pass an empty string; the tool will then create a new thread and return its id. For subsequent calls, pass the previously returned threadId to continue the same conversation.',
    ),
});

const outputSchema = z.object({
  success: z.boolean().describe('Whether the invocation completed successfully'),
  response: z.string().describe('The assistant response text from Mastra Code'),
  threadId: z
    .string()
    .optional()
    .describe(
      'The Mastra Code thread ID used for this invocation (existing or newly created).',
    ),
  error: z.string().optional().describe('Error message if success is false'),
});

function extractTextFromMessage(message: HarnessMessage): string {
  const parts: string[] = [];
  for (const block of message.content) {
    if (block.type === 'text') {
      parts.push(block.text);
    }
  }
  return parts.join('');
}

export const mastraCodeApiTool = createTool({
  id: 'mastra-code-api',
  description:
    'Invokes Mastra Code programmatically for a project: sends a question or task and returns the assistant response. Use the project root path from repositoryListTool (directory field).',
  inputSchema,
  outputSchema,
  execute: async (input) => {
    const { projectPath, query, threadId } = input;

    const { harness } = await createMastraCode({
      cwd: projectPath,
      initialState: { yolo: true },
    });
    await harness.init();
    await harness.switchMode({ modeId: 'plan' });
    await harness.switchModel({ modelId: 'openai/gpt-5.3-codex', scope: 'thread' });

    // Treat blank or whitespace-only threadId as \"no existing thread\"
    let activeThreadId =
      threadId && threadId.trim().length > 0 ? threadId.trim() : undefined;

    if (activeThreadId) {
      await harness.switchThread({ threadId: activeThreadId });
    } else {
      const thread = await harness.createThread({ title: 'Mastra Code API thread' });
      activeThreadId = thread.id ?? harness.getCurrentThreadId();
    }

    try {
      const responsePromise = new Promise<HarnessMessage | null>((resolve, reject) => {
        let resolved = false;
        harness.subscribe((event) => {
          switch (event.type) {
            case 'message_end':
              if (!resolved && event.message.role === 'assistant') {
                resolved = true;
                resolve(event.message);
              }
              break;
            case 'error':
              if (!resolved) {
                resolved = true;
                reject(event.error);
              }
              break;
            default:
              break;
          }
        });
      });

      await harness.sendMessage({ content: query });

      const lastMessage = await responsePromise;
      const responseText = lastMessage
        ? extractTextFromMessage(lastMessage)
        : '';

      return {
        success: true,
        response: responseText || '(No text response from Mastra Code)',
        threadId: activeThreadId,
      };
    } catch (err) {
      const message = err instanceof Error ? err.message : String(err);
      return {
        success: false,
        response: '',
        threadId: activeThreadId,
        error: message,
      };
    } finally {
      try {
        await harness.stopHeartbeats();
        await harness.destroyWorkspace();
      } catch (teardownErr) {
        console.error('[mastraCodeApiTool] Error tearing down harness:', teardownErr);
      }
    }
  },
});