import { z } from "zod";
import fs from "fs";
import path from "path";
import { sessionCredentials } from "../../auth/credentials.js";
import { withAuthParams } from "../helpers/authSchema.js";
import { fetchProjectInfo } from "../../auth/localClient.js";

export function registerAuthTools(server) {
    server.tool(
        "refresh_acai_token",
        `Refresh the Acai JWT token when it has expired (403 "Token no válido" errors). This re-reads the token from the .acai file on disk. If the token on disk is also expired, it calls the Python server to renew it. Use this tool when any other tool fails with a 403 token error.`,
        withAuthParams({}),
        { readOnlyHint: false, destructiveHint: false },
        async (_args, extra) => {
            try {
                const projectDir = process.env.ACAI_PROJECT_DIR || "";
                const acaiFilePath = projectDir ? path.join(projectDir, ".acai") : "";

                if (!acaiFilePath) {
                    return {
                        content: [{ type: "text", text: JSON.stringify({ success: false, error: "ACAI_PROJECT_DIR not set" }) }],
                        isError: true,
                    };
                }

                // Step 1: Try reading fresh token from .acai (Python server may have already refreshed it)
                let token = "";
                let tokenHash = "";
                let domain = "";
                try {
                    const data = JSON.parse(fs.readFileSync(acaiFilePath, "utf-8"));
                    token = data.token || "";
                    tokenHash = data.tokenHash || "";
                    domain = data.domain || "";
                } catch (e) {
                    return {
                        content: [{ type: "text", text: JSON.stringify({ success: false, error: `Cannot read .acai: ${e.message}` }) }],
                        isError: true,
                    };
                }

                // Step 2: Check if token is expired by decoding JWT
                let isExpired = false;
                try {
                    const payload = token.split(".")[1];
                    const decoded = JSON.parse(Buffer.from(payload, "base64").toString());
                    isExpired = Date.now() / 1000 > (decoded.exp || 0) - 300;
                } catch {
                    isExpired = true;
                }

                // Step 3: If expired, ask Python server to refresh it
                if (isExpired) {
                    try {
                        const info = await fetchProjectInfo({ project_dir: projectDir });
                        token = info?.token || token;
                        tokenHash = info?.tokenHash || tokenHash;
                        domain = info?.domain || domain;
                    } catch (e) {
                        return {
                            content: [{ type: "text", text: JSON.stringify({ success: false, error: `Token refresh failed: ${e.message}` }) }],
                            isError: true,
                        };
                    }
                }

                // Step 4: Update credentials in memory from the canonical server resolver
                const info = await fetchProjectInfo({ project_dir: projectDir });
                const webUrl = info?.web_url || process.env.ACAI_WEB_URL || "";
                const apiWebUrl = info?.api_web_url || process.env.ACAI_API_WEB_URL || webUrl;
                const website = info?.domain || domain || process.env.ACAI_WEBSITE || "";
                const freshCreds = {
                    token,
                    tokenHash,
                    website,
                    web_url: webUrl,
                    api_web_url: apiWebUrl,
                    forge_host: info?.forge_host || null,
                    profileName: "stdio",
                    role: "developer",
                };
                sessionCredentials.set("_default", freshCreds);
                if (extra?.sessionId) {
                    sessionCredentials.set(extra.sessionId, freshCreds);
                }

                return {
                    content: [{
                        type: "text",
                        text: JSON.stringify({
                            success: true,
                            message: "Token refreshed successfully",
                            expired_before: isExpired,
                            domain: website,
                        }, null, 2),
                    }],
                };
            } catch (error) {
                return {
                    content: [{ type: "text", text: JSON.stringify({ success: false, error: error.message }) }],
                    isError: true,
                };
            }
        }
    );
}