import { parseYAML } from "./yaml";
import { retryFunctional } from "socket-function/src/batching";
import { httpsRequest } from "socket-function/src/https";

const CANNOT_RETRY = "(CANNOT RETRY)";
import { formatNumber } from "socket-function/src/formatting/format";
import { getAPIKey } from "./apiKeys";

export type MessageHistory = {
    role: "system" | "user" | "assistant";
    content: string;
}[];
export type MessageHistory2 = {
    role: "system" | "user" | "assistant";
    content: string | {
        type: "image_url";
        image_url: {
            url: string;
        };
    }[];
}[];

let totalCost = 0;
export function getTotalCost() {
    return totalCost;
}
type OpenRouterOptions = {
    // If not provided, we'll pop up a modal to ask the user for it, Or, if we're on the server, we'll try to read it from the user folder.
    apiKey?: string;
    provider?: {
        sort?: "throughput" | "price" | "latency",
        // https://openrouter.ai/docs/features/provider-routing#ordering-specific-providers
        order?: string[],
    }
    reasoningEffort?: "low" | "medium" | "high",
};
/** IMPORTANT! Make sure to tell the AI to return yaml. */
export async function yamlOpenRouterCall<T>(config: {
    model: string;
    messages: MessageHistory;
    retries?: number;
    options?: OpenRouterOptions;
    onCost?: (cost: number) => void;
    validate?: (response: T) => void;
}): Promise<T> {
    let { model, messages, retries = 3, options, onCost } = config;
    try {
        let response = await openRouterCall({ model, messages, options, onCost, retries: 0 });
        let result = parseYAML(response) as T;
        config.validate?.(result);
        return result;
    } catch (error) {
        if (retries > 0) {
            return yamlOpenRouterCall<T>({ model, messages, retries: retries - 1, onCost });
        }
        throw error;
    }
}

export async function simpleAICall(model: string, message: string): Promise<string> {
    return await openRouterCallBase({
        model,
        messages: [{ role: "user", content: message }],
    });
}
/** The message must request the result to be returned in YAML (we automatically parse this and return an object). */
export async function simpleAICallTyped<T>(model: string, message: string): Promise<T> {
    return await yamlOpenRouterCall({
        model,
        messages: [{ role: "user", content: message }],
    }) as T;
}

let pendingLog: {
    count: number;
    duration: number;
    cost: number;
} | undefined = undefined;

export async function openRouterCall(config: {
    model: string;
    messages: MessageHistory;
    options?: OpenRouterOptions;
    onCost?: (cost: number) => void;
    retries?: number;
}): Promise<string> {
    return await openRouterCallBase(config);
}

export async function openRouterCallBase(config: {
    model: string;
    messages: MessageHistory2;
    options?: OpenRouterOptions;
    onCost?: (cost: number) => void;
    retries?: number;
}): Promise<string> {
    let { model, messages, options, onCost } = config;
    let openrouterKey = options?.apiKey || await getAPIKey("openrouter.json");
    console.log(`Calling ${model} with ${messages.length} messages`);
    let time = Date.now();
    let stillRunning = true;

    void (async () => {
        while (stillRunning) {
            await new Promise(resolve => setTimeout(resolve, 5000));
            if (stillRunning) {
                console.log("OpenRouter call still running...");
            }
        }
    })();

    try {
        return await retryFunctional(async () => {
            // httpsRequest routes through socket-function's DNS cache and throws on any non-2xx (with the response body in the message), which retryFunctional then retries.
            let responseBody = await httpsRequest(
                "https://openrouter.ai/api/v1/chat/completions",
                Buffer.from(JSON.stringify({
                    model,
                    messages,
                    provider: {
                        sort: "throughput",
                    },
                    usage: { include: true },
                    ...options,
                })),
                "POST",
                false,
                {
                    headers: {
                        "Authorization": `Bearer ${openrouterKey}`,
                        "Content-Type": "application/json",
                    },
                }
            );

            let responseObj = JSON.parse(responseBody.toString()) as {
                usage: {
                    cost: number;
                };
                choices: {
                    finish_reason?: string;
                    message: {
                        content: string;
                        refusal?: string | null;
                    };
                }[];
                error?: {
                    code?: number | string;
                    message?: string;
                    metadata?: unknown;
                };
            };
            if (responseObj.error) {
                throw new Error(`OpenRouter returned an error: ${responseObj.error.code} ${responseObj.error.message} ${JSON.stringify(responseObj.error.metadata)}`);
            }
            let choice = responseObj.choices?.[0];
            if (!choice) {
                throw new Error(`OpenRouter returned no choices: ${JSON.stringify(responseObj)}`);
            }
            if (choice.message.refusal) {
                throw new Error(`OpenRouter model refused content: ${choice.message.refusal} ${CANNOT_RETRY}`);
            }
            if (choice.finish_reason === "content_filter") {
                throw new Error(`OpenRouter content filter triggered (finish_reason=content_filter): ${choice.message.content ?? ""} ${CANNOT_RETRY}`);
            }
            if (choice.finish_reason === "error") {
                throw new Error(`OpenRouter completion errored (finish_reason=error): ${choice.message.content ?? ""}`);
            }
            let newCost = responseObj.usage.cost;
            totalCost += newCost;
            onCost?.(newCost);
            if (!pendingLog) {
                pendingLog = {
                    count: 0,
                    duration: 0,
                    cost: 0,
                };
                setTimeout(() => {
                    let log = pendingLog;
                    if (!log) return;
                    console.log(`Ran: ${log.count} calls at a total summed cost of ${formatNumber(1 / log.cost)}/USD`);
                    pendingLog = undefined;
                }, 10_000);
            }
            pendingLog.count++;
            pendingLog.duration += Date.now() - time;
            pendingLog.cost += newCost;
            return responseObj.choices[0].message.content as string;
        }, {
            maxRetries: config.retries || 3,
            shouldRetry: message => !message.includes(CANNOT_RETRY),
        })();
    } finally {
        stillRunning = false;
    }
}