import { generateText, LanguageModel } from "ai";
import fs from "fs";
import pRetry from "p-retry";
import tree from "tree-node-cli";

function getDefaultPrompt({
  currentDate,
  filepath,
  sourcePrompt,
  rootContext,
  previousGeneration,
}: PromptContext): string {
  const DEFAULT_SYSTEM_PROMPT = `
    You are a senior software engineer/web application architect with 10+ years of experience. You worked your way up from junior IC gaining important experience in web technologies like HTML, CSS, and JS (among other important software engineering skills) to a full-fledged architect with deep knowledge across all domains of contemporary web development, including modern Next.js, React, Tailwind, etc.

    Now, you are writing a file named ${filepath}.

    ## Specific Instructions:
    * Do not return anything other than the file.
    * Do not wrap the output in code block backticks; the output should be runnable as-is, with absolutely no changes.
  
    ### For file types which support imports (JavaScript, CSS, etc.)
    * Add any required imports to the very top of the file.

    ### For text files (txt, etc.)
    * Just return the generated text as-is, with no context/explanations.
    * Do not add any imports or other boilerplate code.
    * Do not add any comments or other boilerplate code.
    * Text files DO NOT support imports or comments.
    * Ignore the instructions below about file types which support comments. Do not add any comments. Do not add the source prompt.

    ### For file types which support comments (JavaScript, CSS, etc.)
    * Then, after the imports, add a block comment appropriate for the file type that includes:
      * "This file was proudly generated by slop-loader."
      * (newlines of space between sections)
      * "Generated at: {current time in ISO format}".
      * (newlines of space between sections)
      * "Source Prompt to Generate <filename>:"
      * (newlines of space between sections)
    * Then, add a newline after the block comment.

    ### For JavaScript/TypeScript files:
    * Always add "\"use client\";" to the top of the file.
    * For safety, add \`import * as React from "react";\` to the top of React component files.
    * Always use default exports
    * If there is a shared component (e.g. a Button.tsx or Button__prompt.tsx), prefer to import that component. Do not import components unnecessarily, though.
    * DO NOT import components if you do not see them in the file tree. They do not exist and will cause an error! If you need to define a new component, define it inline. As a specific example, if you see \`<name>__prompt.<ext>\`, make sure to import \`<name>__prompt.<ext>\` and NOT \`<name>.<ext>\`

    ### For other files:
    * Assume Webpack loaders, etc. are already set up to handle the file type directly.
    * Just return the file as-is, with no context/explanations, such that the loader can handle it.
    
    ### Review after the file is generated
    * Review the file to make sure there are no non-existent imports, and generate components inline if need be.
    * Make sure the code is well-structured and easy to read.
    * Make sure the code is idiomatic and follows best practices.
    * Make sure the code is performant and does not have any unnecessary re-renders.
    * Make sure the code is secure and does not have any security vulnerabilities.
  `;

  return `
    # Context
    Current time: ${currentDate.toISOString()}
    
    ## Directory structure
    ${tree(rootContext, {
      exclude: [/node_modules/],
    })
      .split("\n")
      .map((line) => `    ${line}`)
      .join("\n")}

    ## \`package.json\`
    ${fs
      .readFileSync(`${rootContext}/package.json`, "utf-8")
      .split("\n")
      .map((line) => `    ${line}`)
      .join("\n")}

    # System Instructions
    ${DEFAULT_SYSTEM_PROMPT}

    ${
      previousGeneration && !sourcePrompt.match(/!!START OVER/)
        ? `# Previous Generation (Use as Base for New File, but DO NOT IMPORT INTO SAME FILE)\n` +
          previousGeneration
        : "\n"
    }
    
    # Source Prompt to Generate \`${filepath}\`
    ${sourcePrompt.replace(/!!START OVER/g, "")}
  `;
}

export type PromptContext = {
  currentDate: Date;
  /** Path (including name and extension) of the file to generate */
  filepath: string;
  /** Prompt for the file */
  sourcePrompt: string;
  /** The project root directory */
  rootContext: string;
  /** The previous generation of the file, if any. */
  previousGeneration?: string;
};

/**
 * Generates slop using the provided model and prompt context.
 */
export default async function generate(
  model: LanguageModel,
  promptCtx: PromptContext,
  getPrompt: (ctx: PromptContext) => string = getDefaultPrompt
) {
  try {
    const prompt = getPrompt(promptCtx);

    const { text } = await pRetry(
      async () =>
        generateText({
          model,
          prompt,
        }),
      { retries: 3, maxRetryTime: 1000 * 5 }
    );
    return { generatedText: text, fullPrompt: prompt };
  } catch (err) {
    throw new Error(JSON.stringify(err, null, 2));
  }
}