/***
  @ Base: https://play.google.com/store/apps/details?id=com.mediagenerator.toonmix
  @ Author: Shannz
  @ Note: Anime art generator with ai.
***/

import axios from 'axios';
import fs from 'fs';
import crypto from 'crypto';

const CONFIG = {
    AUTH_URL: "https://www.googleapis.com/identitytoolkit/v3/relyingparty/signupNewUser?key=AIzaSyCtEPD7MFc27-uR8FB2diwOy9FxXOsHu1o",
    API_URL: "https://us-central1-media-generator-b03ac.cloudfunctions.net",
    HEADERS: {
        COMMON: {
            'User-Agent': 'okhttp/3.12.13',
            'Accept-Encoding': 'gzip',
            'Content-Type': 'application/json; charset=utf-8',
            'firebase-instance-id-token': 'e4u8BWbsR0Kmsh1mJp7eLW:APA91bESjXNGuiWgSc_Wj7xBpMw4FIQDwT3ZcsCfooTg5Yw20cIHlWR4Ogwjz-Z_uG2uLcKT0BGh9agoehsxZOI_JtVkn6fitN_LzYtksnJDzAYb-HZvBRk'
        },
        SIGNUP: {
            'User-Agent': 'Dalvik/2.1.0 (Linux; U; Android 15; 25028RN03A Build/AP3A.240905.015.A2)',
            'Connection': 'Keep-Alive',
            'X-Android-Package': 'com.mediagenerator.toonmix',
            'X-Android-Cert': '61ED377E85D386A8DFEE6B864BD85B0BFAA5AF81',
            'Accept-Language': 'id-ID, en-US',
            'X-Client-Version': 'Android/Fallback/X24000001/FirebaseCore-Android',
            'X-Firebase-GMPID': '1:29143983185:android:0af50ff515b110ffd74470',
            'X-Firebase-Client': 'H4sIAAAAAAAA_6tWykhNLCpJSk0sKVayio7VUSpLLSrOzM9TslIyUqoFAFyivEQfAAAA'
        }
    }
};

let session = {
    idToken: null,
    refreshToken: null,
    localId: null,
    expiry: 0
};

const uploadToCloud = async (buffer) => {
    try {
        const filename = `toonmix-${crypto.randomUUID()}.png`;
        const { data } = await axios.post('https://api.cloudsky.biz.id/get-upload-url', {
            fileKey: filename,
            contentType: 'image/png',
            fileSize: buffer.length
        });

        await axios.put(data.uploadUrl, buffer, {
            headers: { 
                'Content-Type': 'image/png', 
                'Content-Length': buffer.length
            }
        });

        return `https://api.cloudsky.biz.id/file?key=${encodeURIComponent(filename)}`;
    } catch (error) {
        console.error(`[Cloud Upload Error]: ${error.message}`);
        return null;
    }
};

const toBase64 = async (input) => {
    try {
        let buffer;
        if (Buffer.isBuffer(input)) {
            buffer = input;
        } else if (input.startsWith('http')) {
            const res = await axios.get(input, { responseType: 'arraybuffer' });
            buffer = Buffer.from(res.data);
        } else if (fs.existsSync(input)) {
            buffer = fs.readFileSync(input);
        } else {
            return null;
        }
        return buffer.toString('base64');
    } catch (e) { return null; }
};

const createPayload = (prompt, params = {}) => {
    const defaultParams = {
        seed: Math.floor(Math.random() * 1000000000),
        width: 1024,
        height: 1024,
        steps: 28,
        cfg_scale: 5,
        model: "nai-diffusion-4-5-curated",
        sampler: "k_euler_ancestral"
    };

    const p = { ...defaultParams, ...params };

    return {
        "negative_prompt": p.negative_prompt || null,
        "params": {
            "cfg_rescale": 0,
            "seed": {
                "@type": "type.googleapis.com/google.protobuf.Int64Value",
                "value": String(p.seed)
            },
            "cfgScale": p.cfg_scale,
            "noise_schedule": "karras",
            "width": {
                "@type": "type.googleapis.com/google.protobuf.Int64Value",
                "value": String(p.width)
            },
            "model": p.model,
            "steps": {
                "@type": "type.googleapis.com/google.protobuf.Int64Value",
                "value": String(p.steps)
            },
            "height": {
                "@type": "type.googleapis.com/google.protobuf.Int64Value",
                "value": String(p.height)
            },
            "cfg_scale": p.cfg_scale,
            "sampler": p.sampler,
            "add_quality_tags": true,
            "negative_quality_tags_type": "heavy"
        },
        "prompt": prompt
    };
};

export const toonmix = {
    signup: async () => {
        try {
            const payload = { "clientType": "CLIENT_TYPE_ANDROID" };
            
            const response = await axios.post(CONFIG.AUTH_URL, payload, {
                headers: CONFIG.HEADERS.SIGNUP
            });

            const data = response.data;
            if (data.idToken) {
                session.idToken = data.idToken;
                session.refreshToken = data.refreshToken;
                session.localId = data.localId;
                session.expiry = Date.now() + (parseInt(data.expiresIn) * 1000);
                return session.idToken;
            }
            return null;
        } catch (error) {
            console.error(`[Signup Error]: ${error.message}`);
            return null;
        }
    },

    generateImage: async (prompt, options = {}) => {
        try {
            if (!session.idToken) await toonmix.signup();
            if (!session.idToken) return { success: false, msg: 'Auth Failed' };

            const payload = {
                data: createPayload(prompt, options)
            };

            const response = await axios.post(`${CONFIG.API_URL}/generateImage`, payload, {
                headers: {
                    ...CONFIG.HEADERS.COMMON,
                    'authorization': `Bearer ${session.idToken}`
                }
            });

            if (response.data?.result?.image) {
                const imgBuffer = Buffer.from(response.data.result.image, 'base64');
                const url = await uploadToCloud(imgBuffer);
                return {
                    success: true,
                    url: url,
                    metadata: response.data.result.metadata
                };
            }

            return { success: false, msg: 'No image returned from API' };

        } catch (error) {
            console.error(`[Generate Error]: ${error.message}`);
            return { success: false, msg: error.message };
        }
    },

    genImageByReference: async (imageReference, prompt, options = {}) => {
        try {
            if (!session.idToken) await toonmix.signup();
            if (!session.idToken) return { success: false, msg: 'Auth Failed' };

            const base64Img = await toBase64(imageReference);
            if (!base64Img) return { success: false, msg: 'Invalid Input Image' };

            const payloadParams = createPayload(prompt, options);

            const payload = {
                data: {
                    ...payloadParams,
                    image: base64Img,
                    strength: options.strength || 0.7,
                    noise: options.noise || 0.2
                }
            };

            const response = await axios.post(`${CONFIG.API_URL}/generateImg2Img`, payload, {
                headers: {
                    ...CONFIG.HEADERS.COMMON,
                    'authorization': `Bearer ${session.idToken}`
                }
            });

            if (response.data?.result?.image) {
                const imgBuffer = Buffer.from(response.data.result.image, 'base64');
                const url = await uploadToCloud(imgBuffer);
                return {
                    success: true,
                    url: url,
                    metadata: response.data.result.metadata
                };
            }

            return { success: false, msg: 'No image returned from API' };

        } catch (error) {
            console.error(`[Img2Img Error]: ${error.message}`);
            return { success: false, msg: error.message };
        }
    }
};