import { z } from "zod";
import axios from "axios";
import fs from "fs";
import path from "path";
import sharp from "sharp";
import { withAuth, getSessionCredentials } from "../../auth/index.js";
import { handleToolError } from "../helpers/errorHandler.js";
import { withAuthParams } from "../helpers/authSchema.js";

// --- Verificación de créditos y reporte de uso ---
const WS_BASE = "https://ws.cocosolution.com/api/handler_acaicode.php";
// Precios Gemini 2.5 Flash: input $0.15/1M tokens, output $0.60/1M tokens
function calcCost(usageMetadata) {
    const input = usageMetadata?.promptTokenCount || 0;
    const output = usageMetadata?.candidatesTokenCount || 0;
    return Math.round(((input * 0.15 + output * 0.60) / 1_000_000) * 1e6) / 1e6;
}

function getAcaiToken() {
    const projectDir = process.env.ACAI_PROJECT_DIR || "";
    if (!projectDir) return null;
    try {
        const acaiFile = path.join(projectDir, ".acai");
        const data = JSON.parse(fs.readFileSync(acaiFile, "utf-8"));
        return data.token || null;
    } catch { return null; }
}

async function checkCredits() {
    const token = getAcaiToken();
    if (!token) return false; // Si no hay token, no bloquear
    const testParam = process.env.STRIPE_MODE === "test" ? "&test" : "";
    try {
        const resp = await axios.put(`${WS_BASE}?action=getUsageLimits${testParam}`, {}, {
            headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json" },
            timeout: 10000,
        });
        return resp.data?.data?.exceeded === true;
    } catch { return false; }
}

function reportImageUsage(usageMetadata, model) {
    const token = getAcaiToken();
    if (!token) return;
    const testParam = process.env.STRIPE_MODE === "test" ? "&test" : "";
    const cost = calcCost(usageMetadata);
    const payload = {
        action: "reportUsage",
        model: model || "gemini-2.5-flash-image",
        input_tokens: usageMetadata?.promptTokenCount || 0,
        output_tokens: usageMetadata?.candidatesTokenCount || 0,
        cache_read_tokens: 0,
        cache_creation_tokens: 0,
        cost_usd: cost,
        session_id: "",
    };
    // Fire and forget
    axios.put(`${WS_BASE}?action=reportUsage${testParam}`, payload, {
        headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json" },
        timeout: 10000,
    }).then(resp => {
        if (resp.data?.success) console.error(`[generate_image] Usage reported: ${model} cost=$${cost}`);
        else console.error(`[generate_image] Usage report failed:`, resp.data);
    }).catch(err => console.error(`[generate_image] Usage report error:`, err.message));
}

export function registerGenerateImageTool(server) {
    server.tool(
        "generate_image",
        `Generate an AI image and save it to the project's uploads folder. Returns preview URLs plus the recommended upload URL for upload_record_image. In Forge environments, prefer uploadUrl (or fullUrl if uploadUrl is absent) over dockerUrl when assigning the image to a record field.`,
        withAuthParams({
            prompt: z.string().describe("Description of the image to generate"),
            width: z.number().optional().describe("Image width in pixels (default: 1024)"),
            height: z.number().optional().describe("Image height in pixels (default: 1024)"),
            style: z.string().optional().describe("Image style hint to add to prompt (e.g., 'photographic', 'digital-art', 'minimalist')"),
            fileName: z.string().optional().describe("Custom filename (without extension). If not provided, auto-generated."),
        }),
        { readOnlyHint: false, destructiveHint: false },
        withAuth(async ({ prompt, width = 1024, height = 1024, style, fileName }, extra) => {
            try {
                const nanoBananaApiKey = process.env.NANO_BANANA_API_KEY;
                if (!nanoBananaApiKey) {
                    return {
                        content: [{ type: "text", text: "Error: NANO_BANANA_API_KEY not set." }],
                        isError: true,
                    };
                }

                const projectDir = process.env.ACAI_PROJECT_DIR || "";
                if (!projectDir) {
                    return {
                        content: [{ type: "text", text: "Error: ACAI_PROJECT_DIR not set." }],
                        isError: true,
                    };
                }

                // Verificar créditos antes de generar
                const exceeded = await checkCredits();
                if (exceeded) {
                    return {
                        content: [{ type: "text", text: "Error: No te quedan créditos. Mejora tu plan para seguir usando el asistente." }],
                        isError: true,
                    };
                }

                // Build prompt with style hint
                const fullPrompt = style ? `${prompt}. Style: ${style}` : prompt;

                // Generate image via Google Gemini
                const geminiModel = process.env.NANO_BANANA_MODEL || "gemini-2.5-flash-image";
                const apiUrl = process.env.NANO_BANANA_URL ||
                    `https://generativelanguage.googleapis.com/v1beta/models/${geminiModel}:generateContent`;

                const generateResponse = await axios.post(
                    apiUrl,
                    {
                        contents: [{ parts: [{ text: fullPrompt }] }],
                        generationConfig: {
                            responseModalities: ["TEXT", "IMAGE"],
                        },
                    },
                    {
                        headers: {
                            "x-goog-api-key": nanoBananaApiKey,
                            "Content-Type": "application/json",
                        },
                        timeout: 120000,
                        validateStatus: (status) => status < 500,
                    }
                );

                // Extract image from response
                let imageBuffer = null;
                if (generateResponse.data.candidates?.[0]?.content?.parts) {
                    for (const part of generateResponse.data.candidates[0].content.parts) {
                        if (part.inlineData?.data) {
                            imageBuffer = Buffer.from(part.inlineData.data, "base64");
                            break;
                        }
                        if (part.text?.startsWith("data:image")) {
                            const match = part.text.match(/data:image\/[^;]+;base64,(.+)/);
                            if (match) {
                                imageBuffer = Buffer.from(match[1], "base64");
                                break;
                            }
                        }
                    }
                }

                if (!imageBuffer) {
                    return {
                        content: [{
                            type: "text",
                            text: `Error: Could not extract image from API response. Status: ${generateResponse.status}. Response: ${JSON.stringify(generateResponse.data).substring(0, 1000)}`
                        }],
                        isError: true,
                    };
                }

                // Compress to JPEG
                const originalSize = imageBuffer.length;
                try {
                    imageBuffer = await sharp(imageBuffer)
                        .jpeg({ quality: 85 })
                        .toBuffer();
                    console.error(`[generate_image] Compressed: ${Math.round(originalSize / 1024)}KB → ${Math.round(imageBuffer.length / 1024)}KB`);
                } catch (e) {
                    console.error(`[generate_image] Compression failed, using original:`, e.message);
                }

                // Save to cms/uploads/generated/
                const uploadsDir = path.join(projectDir, "cms", "uploads", "generated");
                fs.mkdirSync(uploadsDir, { recursive: true });

                const safeName = fileName
                    ? fileName.replace(/[^\w\-]/g, "_") + ".jpg"
                    : `generated-${Date.now()}.jpg`;
                const filePath = path.join(uploadsDir, safeName);
                fs.writeFileSync(filePath, imageBuffer);

                const relativePath = `cms/uploads/generated/${safeName}`;
                const dockerUrl = `http://localhost/${relativePath}`;

                // Reportar uso (fire and forget)
                reportImageUsage(generateResponse.data.usageMetadata, geminiModel);

                const credentials = await getSessionCredentials(extra.sessionId);
                const fullUrl = credentials.web_url ? `${credentials.web_url}/${relativePath}` : dockerUrl;
                const uploadUrl = credentials.web_url ? fullUrl : dockerUrl;

                return {
                    content: [{
                        type: "text",
                        text: JSON.stringify({
                            success: true,
                            prompt: fullPrompt,
                            fileName: safeName,
                            filePath,
                            relativePath,
                            dockerUrl,
                            fullUrl,
                            uploadUrl,
                            size: `${Math.round(imageBuffer.length / 1024)}KB`,
                            note: `Image saved. To assign it with upload_record_image, use imageUrl="${uploadUrl}". dockerUrl is mainly for local preview/debugging.`,
                        }, null, 2),
                    }],
                };
            } catch (error) {
                return handleToolError(error, "generate_image", { prompt });
            }
        })
    );
}