const Anthropic = require('@anthropic-ai/sdk');
const path = require('path');
const YAML = require('yaml');
const fs = require('fs');

// Initialize Anthropic SDK
const anthropic = new Anthropic({
  apiKey: process.env.ANTHROPIC_API_KEY,
});

// Main function
async function main() {
  // Get command line arguments
  const args = process.argv.slice(2);
  const roots = [];
  const includes = [];
  const excludes = [];
  let i = 0;
  while (i < args.length) {
    if (args[i] === '--root') {
      roots.push(args[i + 1]);
      i += 2;
    } else if (args[i] === '--include') {
      includes.push(args[i + 1]);
      i += 2;
    } else if (args[i] === '--exclude') {
      excludes.push(args[i + 1]);
      i += 2;
    } else {
      i++;
    }
  }

  // Get files
  const files = getFiles({
    roots: roots.length > 0 ? roots : ['app'],
    includes: includes.length > 0 ? includes : ['route'],
    excludes: excludes.length > 0 ? excludes : ['dist'],
  });
  console.log('List of files to generate OpenAPI spec for:')
  console.log(files)

  // Initialize OpenAPI spec
  // let spec = initializeSpec();

  // Generate spec for each endpoint
  let yamlDraft = ''
  for (let endpoint of files) {
    const endpointSchema = await generateEndpointSchema(endpoint);
    console.log(endpointSchema)
    yamlDraft += endpointSchema

  }
  fs.writeFileSync('draft_spec.yaml', yamlDraft);

  // Validate and write spec to file
  const validSpec = await validateSpec(yamlDraft);
  writeSpecToFile(validSpec);
}



// Generate endpoint schema
async function generateEndpointSchema(endpoint) {
  const completion = await anthropic.completions.create({
    model: 'claude-2',
    prompt: generatePrompt(endpoint),
    max_tokens_to_sample: 600,
  });
  // get the yaml from the completion

  const regex = /<yaml>(.*?)<\/yaml>/s;
  const match = completion.completion.match(regex);
  return match ? match[1] : null;

}

// Generate prompt for endpoint schema
function generatePrompt(endpoint) {
  return `
Human: Write hello between <yaml></yaml>
Assistant: <yaml>hello</yaml>
Human: Please provide an OpenAPI 3.0 schema for this API endpoint between yaml:
<yaml>
${endpoint}:
${fs.readFileSync(endpoint, 'utf8')}
</yaml>
Assistant:`;
}
// Validate spec
async function validateSpec(spec) {
  const specYaml = YAML.stringify(spec, {
    aliasDuplicateObjects: true, indent: 2,
    lineWidth: -1
  });

  const validYaml = await anthropic.completions.create({
    model: 'claude-2',
    max_tokens_to_sample: 100000,
    prompt: generateValidationPrompt(specYaml)
  });

  const regex = /<yaml>(.*?)<\/yaml>/s;
  const match = validYaml.completion.match(regex);
  return match ? match[1] : null;
}

// Generate validation prompt
function generateValidationPrompt(specYaml) {
  return `
Human: Output valid OpenAPI yaml within <yaml></yaml> always:
Assistant: 
<yaml>
openapi: 3.0.0
info:
  title: Example API
  description: An example to demonstrate OpenAPI 3.0
  version: 1.0.0
servers:
  - url: https://api.example.com
paths:
  /users:
    get:
      summary: Gets a list of users
      responses:
        200:
          description: Success
          content:
            application/json:
              schema:
                type: array
                items:
                  $ref: '#/components/schemas/User'
  /users/{userId}:
    get:
      summary: Gets a user by ID
      parameters:
        - name: userId
          in: path
          required: true
          schema:
            type: integer
      responses:
        200:
          description: Success
          content:
            application/json:
              schema: 
                $ref: '#/components/schemas/User'
        404:
          description: User not found
    put:
      summary: Updates a user
      parameters:
        - name: userId
          in: path
          required: true
          schema:
            type: integer
      requestBody:
        content:
          application/json:
            schema:
              $ref: '#/components/schemas/UserUpdate'
      responses:
        200:
          description: Success
        404:
          description: User not found
 
components:
  schemas:
    User:
      type: object
      properties:
        id:
          type: integer
        name: 
          type: string
    UserUpdate:
      type: object
      properties:
        name:
          type: string
</yaml>
Human:  ${specYaml}
Output valid OpenAPI yaml within <yaml></yaml> fix errors and output a single yaml file:
Assistant:
`;
}

// Write spec to file
function writeSpecToFile(validSpec) {
  if (validSpec) {
    fs.writeFileSync('spec.yaml', validSpec);
  }
}

// Get files
function getFiles(options) {
  const {
    roots = ['./'],
    includes = [],
    excludes = ['/node_modules/'],
    extensions = ['js', 'ts', 'jsx', 'tsx']
  } = options;

  let files = [];
  for (const root of roots) {
    files = files.concat(walk(root));
  }

  return files.filter(file => {
    if (includes.length > 0) {
      return includes.some(include => file.includes(include));
    }
    return !excludes.some(exclude => file.includes(exclude))
  }).filter(file => extensions.includes(path.extname(file).slice(1)));
}

// Walk directory
function walk(dir) {
  let files = [];
  for (const file of fs.readdirSync(dir)) {
    const filepath = path.join(dir, file);
    if (fs.statSync(filepath).isDirectory()) {
      files = files.concat(walk(filepath));
    } else {
      files.push(filepath);
    }
  }
  return files;
}

// Run main function
main();