import * as fs from 'fs/promises';
import * as path from 'path';
import axios from 'axios';
import dotenv from 'dotenv';

dotenv.config();

interface ImageSize {
    width: number;
    height: number;
}

async function convertImage(
    prompt: string,
    baseFileName: string,
    comfyBaseUrl: string,
    comfyOutputDir: string,
    size: ImageSize = { width: 720, height: 1280 },
    useEmpltyLatent: boolean = false
): Promise<string> {
    const COMFY_BASE_URL = comfyBaseUrl.replace(/\/$/, '');
    const COMFY_OUTPUT_DIR = comfyOutputDir;
    let workflow;

    if (useEmpltyLatent) {
        workflow = JSON.parse(await fs.readFile('src/comfyworkflows/edit_image_2_qwen_empty.json', 'utf-8'));
        workflow['21']['inputs']['value'] = prompt;
        workflow['30']['inputs']['width'] = size.width;
        workflow['31']['inputs']['height'] = size.height;
        workflow['14']['inputs']['image'] = baseFileName;
    } else {
        workflow = JSON.parse(await fs.readFile('src/comfyworkflows/edit_image_qwen.json', 'utf-8'));
        workflow['21']['inputs']['value'] = prompt;
        workflow['23']['inputs']['width'] = size.width;
        workflow['23']['inputs']['height'] = size.height;
        workflow['14']['inputs']['image'] = baseFileName;
    }


    const response = await axios.post(`${COMFY_BASE_URL}/prompt`, { prompt: workflow });
    const promptId = response.data.prompt_id;

    let history;
    do {
        await new Promise(resolve => setTimeout(resolve, 1000));
        const historyResponse = await axios.get(`${COMFY_BASE_URL}/history/${promptId}`);
        history = historyResponse.data[promptId];
    } while (!history || Object.keys(history.outputs).length === 0);

    const files = await fs.readdir(COMFY_OUTPUT_DIR!);
    const generatedFiles = files.filter(file => file.startsWith('qwenedit'));

    const fileStats = await Promise.all(
        generatedFiles.map(async (file) => {
            const stat = await fs.stat(path.join(COMFY_OUTPUT_DIR!, file));
            return { file, mtime: stat.mtime };
        })
    );

    fileStats.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());

    const latestFile = fileStats[0].file;
    const newFilePath = path.resolve('./generated', baseFileName);

    await fs.mkdir('./generated', { recursive: true });

    const sourcePath = path.join(COMFY_OUTPUT_DIR!, latestFile);
    try {
        await fs.unlink(newFilePath);
    } catch (err) {
        // ignore if not exists
    }

    await fs.copyFile(sourcePath, newFilePath);

    return newFilePath;
}

// basefilename is connected to image2
// sencondfilename is connect to image1
async function convertImageWithFile(
    prompt: string,
    baseFileName: string,
    secondFileName: string,
    comfyBaseUrl: string,
    comfyOutputDir: string,
    size: ImageSize = { width: 720, height: 1280 }
): Promise<string> {
    const COMFY_BASE_URL = comfyBaseUrl.replace(/\/$/, '');
    const COMFY_OUTPUT_DIR = comfyOutputDir;
    let workflow;

    workflow = JSON.parse(await fs.readFile('src/comfyworkflows/edit_image_2_qwen.json', 'utf-8'));
    workflow['21']['inputs']['value'] = prompt;
    workflow['31']['inputs']['Value'] = size.width;
    workflow['32']['inputs']['Value'] = size.height;
    workflow['14']['inputs']['image'] = baseFileName;
    workflow['29']['inputs']['image'] = secondFileName;

    const response = await axios.post(`${COMFY_BASE_URL}/prompt`, { prompt: workflow });
    const promptId = response.data.prompt_id;

    let history;
    do {
        await new Promise(resolve => setTimeout(resolve, 1000));
        const historyResponse = await axios.get(`${COMFY_BASE_URL}/history/${promptId}`);
        history = historyResponse.data[promptId];
    } while (!history || Object.keys(history.outputs).length === 0);

    const files = await fs.readdir(COMFY_OUTPUT_DIR!);
    const generatedFiles = files.filter(file => file.startsWith('combined'));

    const fileStats = await Promise.all(
        generatedFiles.map(async (file) => {
            const stat = await fs.stat(path.join(COMFY_OUTPUT_DIR!, file));
            return { file, mtime: stat.mtime };
        })
    );

    fileStats.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());

    const latestFile = fileStats[0].file;
    const newFilePath = path.resolve('./generated', baseFileName);

    await fs.mkdir('./generated', { recursive: true });

    const sourcePath = path.join(COMFY_OUTPUT_DIR!, latestFile);
    try {
        await fs.unlink(newFilePath);
    } catch (err) {
        // ignore if not exists
    }

    await fs.copyFile(sourcePath, newFilePath);

    return newFilePath;
}


// basefilename is connected to image1
// sencondfilename is connect to image2
async function convertImageWithFileHandbag(
    prompt: string,
    baseFileName: string,
    secondFileName: string,
    comfyBaseUrl: string,
    comfyOutputDir: string,
    size: ImageSize = { width: 720, height: 1280 }
): Promise<string> {
    const COMFY_BASE_URL = comfyBaseUrl.replace(/\/$/, '');
    const COMFY_OUTPUT_DIR = comfyOutputDir;
    let workflow;

    workflow = JSON.parse(await fs.readFile('src/comfyworkflows/edit_image_2_qwen_handbag.json', 'utf-8'));
    workflow['21']['inputs']['value'] = prompt;
    workflow['31']['inputs']['Value'] = size.width;
    workflow['32']['inputs']['Value'] = size.height;
    workflow['14']['inputs']['image'] = baseFileName;
    workflow['29']['inputs']['image'] = secondFileName;

    const response = await axios.post(`${COMFY_BASE_URL}/prompt`, { prompt: workflow });
    const promptId = response.data.prompt_id;

    let history;
    do {
        await new Promise(resolve => setTimeout(resolve, 1000));
        const historyResponse = await axios.get(`${COMFY_BASE_URL}/history/${promptId}`);
        history = historyResponse.data[promptId];
    } while (!history || Object.keys(history.outputs).length === 0);

    const files = await fs.readdir(COMFY_OUTPUT_DIR!);
    const generatedFiles = files.filter(file => file.startsWith('combined'));

    const fileStats = await Promise.all(
        generatedFiles.map(async (file) => {
            const stat = await fs.stat(path.join(COMFY_OUTPUT_DIR!, file));
            return { file, mtime: stat.mtime };
        })
    );

    fileStats.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());

    const latestFile = fileStats[0].file;
    const newFilePath = path.resolve('./generated', baseFileName);

    await fs.mkdir('./generated', { recursive: true });

    const sourcePath = path.join(COMFY_OUTPUT_DIR!, latestFile);
    try {
        await fs.unlink(newFilePath);
    } catch (err) {
        // ignore if not exists
    }

    await fs.copyFile(sourcePath, newFilePath);

    return newFilePath;
}

// basefilename is connected to image1
// sencondfilename is connect to image2
async function convertImageWithFileForPose(
    prompt: string,
    baseFileName: string,
    secondFileName: string,
    comfyBaseUrl: string,
    comfyOutputDir: string,
    size: ImageSize = { width: 720, height: 1280 }
): Promise<string> {
    const COMFY_BASE_URL = comfyBaseUrl.replace(/\/$/, '');
    const COMFY_OUTPUT_DIR = comfyOutputDir;
    let workflow;

    workflow = JSON.parse(await fs.readFile('src/comfyworkflows/edit_image_2_qwen_pose.json', 'utf-8'));
    workflow['21']['inputs']['value'] = prompt;
    workflow['31']['inputs']['Value'] = size.width;
    workflow['32']['inputs']['Value'] = size.height;
    workflow['14']['inputs']['image'] = baseFileName;
    workflow['29']['inputs']['image'] = secondFileName;

    const response = await axios.post(`${COMFY_BASE_URL}/prompt`, { prompt: workflow });
    const promptId = response.data.prompt_id;

    let history;
    do {
        await new Promise(resolve => setTimeout(resolve, 1000));
        const historyResponse = await axios.get(`${COMFY_BASE_URL}/history/${promptId}`);
        history = historyResponse.data[promptId];
    } while (!history || Object.keys(history.outputs).length === 0);

    const files = await fs.readdir(COMFY_OUTPUT_DIR!);
    const generatedFiles = files.filter(file => file.startsWith('combined'));

    const fileStats = await Promise.all(
        generatedFiles.map(async (file) => {
            const stat = await fs.stat(path.join(COMFY_OUTPUT_DIR!, file));
            return { file, mtime: stat.mtime };
        })
    );

    fileStats.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());

    const latestFile = fileStats[0].file;
    const newFilePath = path.resolve('./generated', baseFileName);

    await fs.mkdir('./generated', { recursive: true });

    const sourcePath = path.join(COMFY_OUTPUT_DIR!, latestFile);
    try {
        await fs.unlink(newFilePath);
    } catch (err) {
        // ignore if not exists
    }

    await fs.copyFile(sourcePath, newFilePath);

    return newFilePath;
}


export async function extractCloth(
    prompt: string,
    baseFileName: string,
    comfyBaseUrl: string,
    comfyOutputDir: string,
    size: ImageSize = { width: 720, height: 1280 }
): Promise<string> {
    const COMFY_BASE_URL = comfyBaseUrl.replace(/\/$/, '');
    const COMFY_OUTPUT_DIR = comfyOutputDir;
    let workflow;

    workflow = JSON.parse(await fs.readFile('src/comfyworkflows/cloth_extractor.json', 'utf-8'));
    workflow['21']['inputs']['value'] = prompt;
    workflow['24']['inputs']['width'] = size.width;
    workflow['24']['inputs']['height'] = size.height;
    workflow['64']['inputs']['image'] = baseFileName;

    const response = await axios.post(`${COMFY_BASE_URL}/prompt`, { prompt: workflow });
    const promptId = response.data.prompt_id;

    let history;
    do {
        await new Promise(resolve => setTimeout(resolve, 1000));
        const historyResponse = await axios.get(`${COMFY_BASE_URL}/history/${promptId}`);
        history = historyResponse.data[promptId];
    } while (!history || Object.keys(history.outputs).length === 0);

    const files = await fs.readdir(COMFY_OUTPUT_DIR!);
    const generatedFiles = files.filter(file => file.startsWith('qwenedit'));

    const fileStats = await Promise.all(
        generatedFiles.map(async (file) => {
            const stat = await fs.stat(path.join(COMFY_OUTPUT_DIR!, file));
            return { file, mtime: stat.mtime };
        })
    );

    fileStats.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());

    const latestFile = fileStats[0].file;
    const newFilePath = path.resolve('./generated', baseFileName);

    await fs.mkdir('./generated', { recursive: true });

    const sourcePath = path.join(COMFY_OUTPUT_DIR!, latestFile);
    try {
        await fs.unlink(newFilePath);
    } catch (err) {
        // ignore if not exists
    }

    await fs.copyFile(sourcePath, newFilePath);

    return newFilePath;
}



export async function convertImageWithMultipleFile(
    prompt: string,
    srcFiles: string[],
    outputFile: string,
    comfyBaseUrl: string,
    comfyOutputDir: string,
    size: ImageSize = { width: 720, height: 1280 }
): Promise<string> {
    const COMFY_BASE_URL = comfyBaseUrl.replace(/\/$/, '');
    const COMFY_OUTPUT_DIR = comfyOutputDir;
    let workflow;

    workflow = JSON.parse(await fs.readFile('src/comfyworkflows/edit_image_multiple_qwen.json', 'utf-8'));
    workflow['21']['inputs']['value'] = prompt;
    workflow['76']['inputs']['number'] = size.width;
    workflow['77']['inputs']['number'] = size.height;

    if (srcFiles[0])
        workflow['64']['inputs']['image'] = srcFiles[0];

    if (srcFiles[1])
        workflow['15']['inputs']['image'] = srcFiles[1];

    if (srcFiles[2])
        workflow['68']['inputs']['image'] = srcFiles[2];

    const response = await axios.post(`${COMFY_BASE_URL}/prompt`, { prompt: workflow });
    const promptId = response.data.prompt_id;

    let history;
    do {
        await new Promise(resolve => setTimeout(resolve, 1000));
        const historyResponse = await axios.get(`${COMFY_BASE_URL}/history/${promptId}`);
        history = historyResponse.data[promptId];
    } while (!history || Object.keys(history.outputs).length === 0);

    const files = await fs.readdir(COMFY_OUTPUT_DIR!);
    const generatedFiles = files.filter(file => file.startsWith('qwenedit'));

    const fileStats = await Promise.all(
        generatedFiles.map(async (file) => {
            const stat = await fs.stat(path.join(COMFY_OUTPUT_DIR!, file));
            return { file, mtime: stat.mtime };
        })
    );

    fileStats.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());

    const latestFile = fileStats[0].file;
    const newFilePath = path.resolve('./generated', outputFile);

    await fs.mkdir('./generated', { recursive: true });

    const sourcePath = path.join(COMFY_OUTPUT_DIR!, latestFile);
    try {
        await fs.unlink(newFilePath);
    } catch (err) {
        // ignore if not exists
    }

    await fs.copyFile(sourcePath, newFilePath);

    return newFilePath;
}


export async function convertImageVton(
    personFile: string,
    clothFile: string,
    outputFile: string,
    comfyBaseUrl: string,
    comfyOutputDir: string,
    size: ImageSize = { width: 720, height: 1280 }
): Promise<string> {
    const COMFY_BASE_URL = comfyBaseUrl.replace(/\/$/, '');
    const COMFY_OUTPUT_DIR = comfyOutputDir;
    let workflow;

    workflow = JSON.parse(await fs.readFile('src/comfyworkflows/vton_cloth.json', 'utf-8'));
    workflow['76']['inputs']['number'] = size.width;
    workflow['77']['inputs']['number'] = size.height;

    workflow['15']['inputs']['image'] = personFile;
    workflow['64']['inputs']['image'] = clothFile;

    const response = await axios.post(`${COMFY_BASE_URL}/prompt`, { prompt: workflow });
    const promptId = response.data.prompt_id;

    let history;
    do {
        await new Promise(resolve => setTimeout(resolve, 1000));
        const historyResponse = await axios.get(`${COMFY_BASE_URL}/history/${promptId}`);
        history = historyResponse.data[promptId];
    } while (!history || Object.keys(history.outputs).length === 0);

    const files = await fs.readdir(COMFY_OUTPUT_DIR!);
    const generatedFiles = files.filter(file => file.startsWith('qwenedit'));

    const fileStats = await Promise.all(
        generatedFiles.map(async (file) => {
            const stat = await fs.stat(path.join(COMFY_OUTPUT_DIR!, file));
            return { file, mtime: stat.mtime };
        })
    );

    fileStats.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());

    const latestFile = fileStats[0].file;
    const newFilePath = path.resolve('./generated', outputFile);

    await fs.mkdir('./generated', { recursive: true });

    const sourcePath = path.join(COMFY_OUTPUT_DIR!, latestFile);
    try {
        await fs.unlink(newFilePath);
    } catch (err) {
        // ignore if not exists
    }

    await fs.copyFile(sourcePath, newFilePath);

    return newFilePath;
}


export async function convertImageVtonPose(
    personFile: string,
    clothFile: string,
    poseFile: string,
    outputFile: string,
    comfyBaseUrl: string,
    comfyOutputDir: string,
    size: ImageSize = { width: 720, height: 1280 }
): Promise<string> {
    const COMFY_BASE_URL = comfyBaseUrl.replace(/\/$/, '');
    const COMFY_OUTPUT_DIR = comfyOutputDir;
    let workflow;

    workflow = JSON.parse(await fs.readFile('src/comfyworkflows/vton.json', 'utf-8'));
    workflow['76']['inputs']['number'] = size.width;
    workflow['77']['inputs']['number'] = size.height;

    workflow['15']['inputs']['image'] = personFile;
    workflow['64']['inputs']['image'] = clothFile;
    workflow['68']['inputs']['image'] = poseFile;

    const response = await axios.post(`${COMFY_BASE_URL}/prompt`, { prompt: workflow });
    const promptId = response.data.prompt_id;

    let history;
    do {
        await new Promise(resolve => setTimeout(resolve, 1000));
        const historyResponse = await axios.get(`${COMFY_BASE_URL}/history/${promptId}`);
        history = historyResponse.data[promptId];
    } while (!history || Object.keys(history.outputs).length === 0);

    const files = await fs.readdir(COMFY_OUTPUT_DIR!);
    const generatedFiles = files.filter(file => file.startsWith('qwenedit'));

    const fileStats = await Promise.all(
        generatedFiles.map(async (file) => {
            const stat = await fs.stat(path.join(COMFY_OUTPUT_DIR!, file));
            return { file, mtime: stat.mtime };
        })
    );

    fileStats.sort((a, b) => b.mtime.getTime() - a.mtime.getTime());

    const latestFile = fileStats[0].file;
    const newFilePath = path.resolve('./generated', outputFile);

    await fs.mkdir('./generated', { recursive: true });

    const sourcePath = path.join(COMFY_OUTPUT_DIR!, latestFile);
    try {
        await fs.unlink(newFilePath);
    } catch (err) {
        // ignore if not exists
    }

    await fs.copyFile(sourcePath, newFilePath);

    return newFilePath;
}

export { convertImage, convertImageWithFile, convertImageWithFileForPose, convertImageWithFileHandbag };
