const { app, BrowserWindow, ipcMain, protocol, dialog } = require('electron');
const { BedrockRuntimeClient, InvokeModelCommand, InvokeModelWithResponseStreamCommand } = require("@aws-sdk/client-bedrock-runtime");
const { fromIni } = require("@aws-sdk/credential-provider-ini");
const fs = require('fs').promises;
const path = require('path');
const url = require('url');
const axios = require('axios');
const cheerio = require('cheerio');

let bedrockClient;

async function initializeBedrockClient() {
  bedrockClient = new BedrockRuntimeClient({ 
    region: "us-east-1",
    credentials: fromIni({ profile: 'default' })
  });
}

function createWindow() {
  const win = new BrowserWindow({
    width: 800,
    height: 600,
    webPreferences: {
      nodeIntegration: true,
      contextIsolation: false,
    },
  });

  if (process.env.NODE_ENV === 'development') {
    win.loadURL('http://localhost:5173');
  } else {
    win.loadFile(path.join(__dirname, 'build', 'index.html'));
  }

  if (process.env.NODE_ENV === 'development') {
    win.webContents.openDevTools();
  }
}

app.whenReady().then(async () => {
  await initializeBedrockClient();
  
  // Register protocol
  protocol.registerFileProtocol('safe-file', (request, callback) => {
    const filePath = url.fileURLToPath('file://' + request.url.slice('safe-file://'.length));
    callback(filePath);
  });

  
  createWindow();
});

ipcMain.handle('fetch-url-content', async (event, url) => {
    try {
      const response = await axios.get(url);
      const $ = cheerio.load(response.data);
      
      // Extract text content from the page
      const content = $('body').text().trim();
      
      // Limit content to a reasonable length (e.g., 1000 characters)
      const limitedContent = content.substring(0, 1000);
      
      event.sender.send('url-content-fetched', limitedContent);
    } catch (error) {
      console.error('Error fetching URL:', error);
      event.sender.send('url-fetch-error', error.message);
    }
  });
  
  async function readDirectory(dirPath) {
    const fileTree = [];
    const files = await fs.readdir(dirPath, { withFileTypes: true });
  
    for (const file of files) {
      const filePath = path.join(dirPath, file.name);
      if (file.isDirectory()) {
        const subDir = await readDirectory(filePath);
        fileTree.push({ name: file.name, type: 'directory', children: subDir });
      } else {
        let content = '';
        if (file.name.match(/\.(txt|md|js|py|html|css|json)$/i)) {
          content = await fs.readFile(filePath, 'utf-8');
        }
        fileTree.push({ name: file.name, type: 'file', content });
      }
    }
  
    return fileTree;
  }
  
  function formatFileTree(fileTree, indent = '') {
    let result = '';
    for (const item of fileTree) {
      result += `${indent}${item.type === 'directory' ? '📁' : '📄'} ${item.name}\n`;
      if (item.type === 'directory' && item.children) {
        result += formatFileTree(item.children, indent + '  ');
      }
    }
    return result;
  }
  
  ipcMain.handle('select-directory', async (event) => {
    const result = await dialog.showOpenDialog({ properties: ['openDirectory'] });
    if (result.canceled) {
      return null;
    }
    const dirPath = result.filePaths[0];
    const fileTree = await readDirectory(dirPath);
    const formattedTree = formatFileTree(fileTree);
    const fileContents = JSON.stringify(fileTree, null, 2);
    return { tree: formattedTree, contents: fileContents };
  });

  
  ipcMain.handle('generate-response-stream', async (event, { prompt, imagePath, model, temperature, maxTokens, context, fileTree, fileContents }) => {
    if (!bedrockClient) {
      throw new Error('Bedrock client not initialized');
    }
  
    console.log("DEBUG: Generating response for prompt:", prompt);
  
    try {
      const requestBody = {
        anthropic_version: "bedrock-2023-05-31",
        max_tokens: maxTokens,
        temperature: temperature,
        messages: [
          {
            role: "user",
            content: [
              { type: "text", text: `Context: ${context || ''}\n\nFile Tree:\n${fileTree || ''}\n\nFile Contents:\n${fileContents || ''}\n\nUser: ${prompt}` }
            ]
          }
        ]
      };
    
      if (imagePath) {
        const imageBuffer = await fs.readFile(imagePath);
        const base64Image = imageBuffer.toString('base64');
  
        requestBody.messages[0].content.unshift({
          type: "image",
          source: {
            type: "base64",
            media_type: "image/jpeg",
            data: base64Image
          }
        });
      }
  
      const params = {
        modelId: "anthropic.claude-3-sonnet-20240229-v1:0", // or your preferred Claude 3 model
        contentType: "application/json",
        accept: "application/json",
        body: JSON.stringify(requestBody),
      };
  
      const command = new InvokeModelWithResponseStreamCommand(params);
      const response = await bedrockClient.send(command);
  
      for await (const chunk of response.body) {
        if (chunk.chunk && chunk.chunk.bytes) {
          const decodedChunk = JSON.parse(new TextDecoder().decode(chunk.chunk.bytes));
          if (decodedChunk.type === 'content_block_delta' && decodedChunk.delta && decodedChunk.delta.text) {
            event.sender.send('response-chunk', decodedChunk.delta.text);
          }
        }
      }
      
      event.sender.send('response-end');
    } catch (error) {
      console.error("ERROR: Failed to generate text with Bedrock:", error);
      event.sender.send('response-error', error.message);
    }
  });
  
  
ipcMain.handle('generate-response', async (event, { prompt, imagePath }) => {
  if (!bedrockClient) {
    throw new Error('Bedrock client not initialized');
  }

  console.log("DEBUG: Generating response for prompt:", prompt);

  try {
    const requestBody = {
      anthropic_version: "bedrock-2023-05-31",
      max_tokens: 20000,
      temperature: 0.7,
      top_p: 0.9,
      messages: [
        {
          role: "user",
          content: []
        }
      ]
    };

    if (imagePath) {
      const imageBuffer = await fs.readFile(imagePath);
      const base64Image = imageBuffer.toString('base64');

      requestBody.messages[0].content.push({
        type: "image",
        source: {
          type: "base64",
          media_type: "image/jpeg",
          data: base64Image
        }
      });
    }

    requestBody.messages[0].content.push({
      type: "text",
      text: prompt
    });

    const params = {
      modelId: "anthropic.claude-3-sonnet-20240229-v1:0",
      contentType: "application/json",
      accept: "application/json",
      body: JSON.stringify(requestBody),
    };

    const command = new InvokeModelCommand(params);
    const response = await bedrockClient.send(command);
    const responseBody = JSON.parse(new TextDecoder().decode(response.body));
    console.log("DEBUG: Raw response from Bedrock:", JSON.stringify(responseBody));

    if (responseBody.content && Array.isArray(responseBody.content) && responseBody.content.length > 0) {
      const generatedText = responseBody.content[0].text;
      if (generatedText) {
        console.log("DEBUG: Generated text:", generatedText);
        return generatedText;
      } else {
        console.log("WARNING: Generated text is empty. Full response:", JSON.stringify(responseBody));
        return "No response generated";
      }
    } else {
      console.log("WARNING: Unexpected response format. Full response:", JSON.stringify(responseBody));
      return "Unexpected response format";
    }
  } catch (error) {
    console.error("ERROR: Failed to generate text with Bedrock:", error);
    return "Error: " + error.message;
  }
});

ipcMain.handle('generate-image', async (event, { prompt }) => {
    if (!bedrockClient) {
      throw new Error('Bedrock client not initialized');
    }
  
    console.log("DEBUG: Generating image for prompt:", prompt);
  
    try {
      const requestBody = {
        text_prompts: [{ text: prompt }],
        cfg_scale: 10,
        steps: 50,
        seed: Math.floor(Math.random() * 1000000),
      };
  
      const params = {
        modelId: "stability.stable-diffusion-xl-v1",
        contentType: "application/json",
        accept: "application/json",
        body: JSON.stringify(requestBody),
      };
  
      const command = new InvokeModelCommand(params);
      const response = await bedrockClient.send(command);
      const responseBody = JSON.parse(new TextDecoder().decode(response.body));
  
      if (responseBody.artifacts && responseBody.artifacts.length > 0) {
        const imageBase64 = responseBody.artifacts[0].base64;
        const imagePath = path.join(app.getPath('userData'), `generated_image_${Date.now()}.png`);
        await fs.writeFile(imagePath, imageBase64, 'base64');
        return `safe-file://${imagePath}`; // Return the safe-file URL instead of the direct file path
      } else {
        throw new Error('No image generated');
      }    
    } catch (error) {
      console.error("ERROR: Failed to generate image with Bedrock:", error);
      throw error;
    }
  });
  
app.on('window-all-closed', () => {
  if (process.platform !== 'darwin') {
    app.quit();
  }
});

app.on('activate', () => {
  if (BrowserWindow.getAllWindows().length === 0) {
    createWindow();
  }
});