import { HeliconeManualLogger } from "@helicone/helpers";
import OpenAI from "openai";
import fs from "fs";
import dotenv from "dotenv";
dotenv.config();
// Initialize Helicone Manual Logger
const heliconeLogger = new HeliconeManualLogger({
  apiKey: process.env.HELICONE_API_KEY!,
  loggingEndpoint: "https://api.worker.helicone.ai/oai/v1/log",
  headers: {}
});
// Initialize OpenAI client
const openai = new OpenAI({
  apiKey: process.env.OPENAI_API_KEY!,
});
function createBatchFile(filename: string = "data.jsonl") {
  const batchRequests = [
    {
      custom_id: "req-1",
      method: "POST",
      url: "/v1/chat/completions",
      body: {
        model: "gpt-4o-mini",
        messages: [{ 
          role: "user", 
          content: "Write a professional email to schedule a meeting with a client about quarterly business review" 
        }],
        max_tokens: 300
      }
    },
    {
      custom_id: "req-2", 
      method: "POST",
      url: "/v1/chat/completions",
      body: {
        model: "gpt-4o-mini",
        messages: [{ 
          role: "user", 
          content: "Explain the benefits of cloud computing for small businesses in simple terms" 
        }],
        max_tokens: 250
      }
    },
    {
      custom_id: "req-3",
      method: "POST", 
      url: "/v1/chat/completions",
      body: {
        model: "gpt-4o-mini",
        messages: [{ 
          role: "user", 
          content: "Create a Python function that calculates compound interest with proper error handling" 
        }],
        max_tokens: 400
      }
    }
  ];
  const jsonlContent = batchRequests.map(req => JSON.stringify(req)).join('\n');
  fs.writeFileSync(filename, jsonlContent);
  console.log(`Created batch file: ${filename}`);
  return filename;
}
async function uploadFile(filename: string) {
  console.log("Uploading file...");
  
  try {
    const file = await openai.files.create({
      file: fs.createReadStream(filename),
      purpose: "batch",
    });
    
    console.log(`File uploaded: ${file.id}`);
    return file.id;
  } catch (error) {
    console.error("Error uploading file:", error);
    throw error;
  }
}
async function createBatch(fileId: string) {
  console.log("Creating batch...");
  
  try {
    const batch = await openai.batches.create({
      input_file_id: fileId,
      endpoint: "/v1/chat/completions", 
      completion_window: "24h"
    });
    
    console.log(`Batch created: ${batch.id}`);
    console.log(`Status: ${batch.status}`);
    return batch;
  } catch (error) {
    console.error("Error creating batch:", error);
    throw error;
  }
}
async function waitForCompletion(batchId: string) {
  console.log("Waiting for batch completion...");
  
  while (true) {
    try {
      const batch = await openai.batches.retrieve(batchId);
      console.log(`Status: ${batch.status}`);
      
      if (batch.status === "completed") {
        console.log("Batch completed!");
        return batch;
      } else if (batch.status === "failed" || batch.status === "expired" || batch.status === "cancelled") {
        throw new Error(`Batch failed with status: ${batch.status}`);
      }
      
      console.log("Waiting 5 seconds...");
      await new Promise(resolve => setTimeout(resolve, 5000));
    } catch (error) {
      console.error("Error checking batch status:", error);
      throw error;
    }
  }
}
async function retrieveAndLogResults(batch: any) {
  if (!batch.output_file_id || !batch.input_file_id) {
    throw new Error("No output or input file available");
  }
  console.log("Retrieving batch results...");
  
  try {
    // Get original requests
    const inputFileContent = await openai.files.content(batch.input_file_id);
    const inputContent = await inputFileContent.text();
    const originalRequests = inputContent.trim().split('\n').map(line => JSON.parse(line));
    
    // Get batch results
    const outputFileContent = await openai.files.content(batch.output_file_id);
    const outputContent = await outputFileContent.text();
    const results = outputContent.trim().split('\n').map(line => JSON.parse(line));
    
    console.log(`Found ${results.length} results`);
    
    // Create mapping of custom_id to original request
    const requestMap = new Map();
    originalRequests.forEach(req => {
      requestMap.set(req.custom_id, req.body);
    });
    
    // Log each result to Helicone
    for (const result of results) {
      const { custom_id, response } = result;
      
      if (response && response.body) {
        console.log(`\nLogging ${custom_id}...`);
        
        const originalRequest = requestMap.get(custom_id);
        
        if (originalRequest) {
          // Modify model name to distinguish batch requests
          const modifiedRequest = {
            ...originalRequest,
            model: originalRequest.model + "-batch"
          };
          
          const modifiedResponse = {
            ...response.body,
            model: response.body.model + "-batch"
          };
          
          // Log to Helicone with additional metadata
          await heliconeLogger.logSingleRequest(
            modifiedRequest,
            JSON.stringify(modifiedResponse),
            {
              additionalHeaders: {
                "Helicone-User-Id": "batch-demo",
                "Helicone-Property-CustomId": custom_id,
                "Helicone-Property-BatchId": batch.id,
                "Helicone-Property-ProcessingType": "batch",
                "Helicone-Property-Provider": "openai"
              }
            }
          );
          
          const responseText = response.body.choices?.[0]?.message?.content || "No response";
          console.log(`${custom_id}: "${responseText.substring(0, 100)}..."`);
        } else {
          console.log(`Could not find original request for ${custom_id}`);
        }
      }
    }
    
    console.log(`\nSuccessfully logged all ${results.length} requests to Helicone!`);
    return results;
    
  } catch (error) {
    console.error("Error retrieving results:", error);
    throw error;
  }
}
async function main() {
  console.log("OpenAI Batch API with Helicone Logging\n");
  
  // Validate environment variables
  if (!process.env.HELICONE_API_KEY) {
    console.error("Please set HELICONE_API_KEY environment variable");
    return;
  }
  
  if (!process.env.OPENAI_API_KEY) {
    console.error("Please set OPENAI_API_KEY environment variable");
    return;
  }
  try {
    // Complete batch workflow
    const filename = createBatchFile();
    const fileId = await uploadFile(filename);
    const batch = await createBatch(fileId);
    const completedBatch = await waitForCompletion(batch.id);
    await retrieveAndLogResults(completedBatch);
    
    // Cleanup
    if (fs.existsSync(filename)) {
      fs.unlinkSync(filename);
      console.log(`Cleaned up ${filename}`);
    }
    
  } catch (error) {
    console.error("Error:", error);
  }
}
if (require.main === module) {
  main();
}