import { calculateCost, createAssistantMessageEventStream, } from "@earendil-works/pi-ai";
import { readFileSync, existsSync } from "node:fs";
import { resolve } from "node:path";
import { homedir } from "node:os";
import { DefaultAzureCredential } from "@azure/identity";
// =============================================================================
// Token Provider
// =============================================================================
/** Azure AI Foundry scope for Entra ID tokens */
const AZURE_AI_SCOPE = "https://ai.azure.com/.default";
/** Cached token — refreshed when within 5 min of expiry */
let cachedToken = null;
async function getIdentityToken() {
    const now = Date.now();
    const expiryBuffer = 5 * 60 * 1000; // 5 minutes
    if (cachedToken && cachedToken.expiresOnTimestamp - now > expiryBuffer) {
        return cachedToken.token;
    }
    const credential = new DefaultAzureCredential();
    cachedToken = await credential.getToken(AZURE_AI_SCOPE);
    if (!cachedToken)
        throw new Error("[Azure Foundry] Failed to acquire identity token");
    return cachedToken.token;
}
/**
 * Returns a token getter function appropriate for the configured auth type.
 * For api-key: always returns the static key.
 * For azure-identity: fetches/caches an Entra ID token via DefaultAzureCredential.
 */
function makeTokenGetter(auth) {
    if (auth.type === "api-key") {
        return () => Promise.resolve(auth.apiKey);
    }
    return getIdentityToken;
}
function loadConfig() {
    // Search order: project root → ~/.pi/azure-foundry.config.json
    const candidates = [
        resolve(process.cwd(), "azure-foundry.config.json"),
        resolve(homedir(), ".pi", "azure-foundry.config.json"),
    ];
    for (const p of candidates) {
        if (existsSync(p)) {
            console.log(`[Azure Foundry] Loading config from: ${p}`);
            return JSON.parse(readFileSync(p, "utf-8"));
        }
    }
    throw new Error(`azure-foundry.config.json not found. Checked:\n` +
        candidates.map((p) => `  ${p}`).join("\n") +
        `\n\nCreate one in your project root or at ~/.pi/azure-foundry.config.json`);
}
// =============================================================================
// Deployment → Model mapping
// =============================================================================
const MODEL_DEFAULTS = {
    "claude-sonnet-4-5": { contextWindow: 200000, maxTokens: 16384, reasoning: true, input: ["text", "image"] },
    "claude-sonnet-4-6": { contextWindow: 200000, maxTokens: 16384, reasoning: true, input: ["text", "image"] },
    "claude-haiku-4-5": { contextWindow: 200000, maxTokens: 16384, reasoning: false, input: ["text", "image"] },
    "claude-opus-4-5": { contextWindow: 200000, maxTokens: 32000, reasoning: true, input: ["text", "image"] },
    "gpt-5.4-nano": { contextWindow: 128000, maxTokens: 16384, reasoning: false, input: ["text", "image"] },
    "gpt-4o": { contextWindow: 128000, maxTokens: 4096, reasoning: false, input: ["text", "image"] },
    "gpt-4o-mini": { contextWindow: 128000, maxTokens: 4096, reasoning: false, input: ["text", "image"] },
    "Kimi-K2.5": { contextWindow: 131072, maxTokens: 8192, reasoning: false, input: ["text"] },
    "Kimi-K2.6": { contextWindow: 131072, maxTokens: 8192, reasoning: false, input: ["text"] },
};
const FALLBACK = { contextWindow: 128000, maxTokens: 4096, reasoning: false, input: ["text"] };
/** Publisher stored on the model object so streamSimple can route */
const publisherMap = new Map();
const providerAuthMap = new Map();
function deploymentToModel(d) {
    const defaults = MODEL_DEFAULTS[d.modelName ?? d.name] ?? FALLBACK;
    if (d.modelPublisher)
        publisherMap.set(d.name, d.modelPublisher);
    return {
        id: d.name,
        name: d.modelName ?? d.name,
        reasoning: defaults.reasoning,
        input: defaults.input,
        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
        contextWindow: defaults.contextWindow,
        maxTokens: defaults.maxTokens,
    };
}
// =============================================================================
// SSE Stream Parser
// =============================================================================
async function* parseSSE(reader) {
    const decoder = new TextDecoder();
    let buffer = "";
    while (true) {
        const { done, value } = await reader.read();
        if (done)
            break;
        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split("\n");
        buffer = lines.pop() ?? "";
        for (const line of lines) {
            const trimmed = line.trim();
            if (trimmed.startsWith("data: ")) {
                const data = trimmed.slice(6);
                if (data === "[DONE]")
                    return;
                yield data;
            }
        }
    }
}
// =============================================================================
// OpenAI-format message conversion  (for OpenAI / MoonshotAI / etc.)
// =============================================================================
function toOpenAIMessages(systemPrompt, messages) {
    const out = [];
    if (systemPrompt)
        out.push({ role: "system", content: systemPrompt });
    for (let i = 0; i < messages.length; i++) {
        const msg = messages[i];
        if (msg.role === "user") {
            if (typeof msg.content === "string") {
                out.push({ role: "user", content: msg.content });
            }
            else {
                out.push({ role: "user", content: msg.content.map((c) => c.type === "text" ? { type: "text", text: c.text } :
                        c.type === "image" ? { type: "image_url", image_url: { url: `data:${c.mimeType};base64,${c.data}` } } :
                            { type: "text", text: "" }) });
            }
        }
        else if (msg.role === "assistant") {
            const entry = { role: "assistant" };
            const text = msg.content.filter((b) => b.type === "text").map((b) => b.text).join("\n");
            const tcs = msg.content.filter((b) => b.type === "toolCall").map((b) => ({
                id: b.id, type: "function", function: { name: b.name, arguments: JSON.stringify(b.arguments) },
            }));
            if (text)
                entry.content = text;
            if (tcs.length)
                entry.tool_calls = tcs;
            out.push(entry);
        }
        else if (msg.role === "toolResult") {
            const m = msg;
            out.push({ role: "tool", tool_call_id: m.toolCallId, content: m.content.filter((c) => c.type === "text").map((c) => c.text).join("\n") });
        }
    }
    return out;
}
function toOpenAITools(tools) {
    return tools.map((t) => ({ type: "function", function: { name: t.name, description: t.description, parameters: t.parameters } }));
}
// =============================================================================
// Anthropic-format message conversion
// =============================================================================
function toAnthropicMessages(messages) {
    const out = [];
    for (let i = 0; i < messages.length; i++) {
        const msg = messages[i];
        if (msg.role === "user") {
            if (typeof msg.content === "string") {
                out.push({ role: "user", content: msg.content });
            }
            else {
                out.push({ role: "user", content: msg.content.map((c) => c.type === "text" ? { type: "text", text: c.text } :
                        c.type === "image" ? { type: "image", source: { type: "base64", media_type: c.mimeType, data: c.data } } :
                            { type: "text", text: "" }) });
            }
        }
        else if (msg.role === "assistant") {
            const blocks = [];
            for (const b of msg.content) {
                if (b.type === "text" && b.text.trim())
                    blocks.push({ type: "text", text: b.text });
                if (b.type === "thinking")
                    blocks.push({ type: "thinking", thinking: b.thinking, signature: b.thinkingSignature ?? "" });
                if (b.type === "toolCall")
                    blocks.push({ type: "tool_use", id: b.id, name: b.name, input: b.arguments });
            }
            if (blocks.length)
                out.push({ role: "assistant", content: blocks });
        }
        else if (msg.role === "toolResult") {
            const m = msg;
            const text = m.content.filter((c) => c.type === "text").map((c) => c.text).join("\n");
            // Anthropic tool results go inside a user message
            out.push({ role: "user", content: [{ type: "tool_result", tool_use_id: m.toolCallId, content: text, is_error: m.isError }] });
        }
    }
    return out;
}
function toAnthropicTools(tools) {
    return tools.map((t) => ({
        name: t.name, description: t.description,
        input_schema: { type: "object", properties: t.parameters.properties ?? {}, required: t.parameters.required ?? [] },
    }));
}
// =============================================================================
// OpenAI-compatible streaming (OpenAI, MoonshotAI, etc.)
// =============================================================================
function streamOpenAI(model, context, options, output, stream, baseHost, auth) {
    return (async () => {
        const url = `${baseHost}/openai/deployments/${model.id}/chat/completions?api-version=2024-10-21`;
        const body = {
            messages: toOpenAIMessages(context.systemPrompt, context.messages),
            max_tokens: options?.maxTokens ?? model.maxTokens,
            stream: true,
            stream_options: { include_usage: true },
        };
        if (context.tools?.length)
            body.tools = toOpenAITools(context.tools);
        const token = await auth.getToken();
        // OpenAI-compat route: api-key auth uses the "api-key" header;
        // Entra ID (azure-identity) uses "Authorization: Bearer".
        const authHeaders = auth.type === "api-key"
            ? { "api-key": token }
            : { "Authorization": `Bearer ${token}` };
        const response = await fetch(url, {
            method: "POST",
            headers: { "Content-Type": "application/json", ...authHeaders },
            body: JSON.stringify(body),
            signal: options?.signal,
        });
        if (!response.ok) {
            const t = await response.text().catch(() => "");
            throw new Error(`Azure Foundry ${response.status}: ${t.slice(0, 500)}`);
        }
        if (!response.body)
            throw new Error("No response body");
        stream.push({ type: "start", partial: output });
        const tcJsonBufs = new Map();
        const tcContentIdx = new Map();
        const reader = response.body.getReader();
        for await (const data of parseSSE(reader)) {
            let chunk;
            try {
                chunk = JSON.parse(data);
            }
            catch {
                continue;
            }
            if (chunk.usage) {
                output.usage.input = chunk.usage.prompt_tokens ?? 0;
                output.usage.output = chunk.usage.completion_tokens ?? 0;
                output.usage.totalTokens = chunk.usage.total_tokens ?? 0;
                calculateCost(model, output.usage);
            }
            const choice = chunk.choices?.[0];
            if (!choice?.delta)
                continue;
            const delta = choice.delta;
            if (typeof delta.content === "string") {
                let idx = output.content.findIndex((b) => b.type === "text");
                if (idx === -1) {
                    output.content.push({ type: "text", text: "" });
                    idx = output.content.length - 1;
                    stream.push({ type: "text_start", contentIndex: idx, partial: output });
                }
                const block = output.content[idx];
                if (block.type === "text") {
                    block.text += delta.content;
                    stream.push({ type: "text_delta", contentIndex: idx, delta: delta.content, partial: output });
                }
            }
            if (delta.tool_calls) {
                for (const tc of delta.tool_calls) {
                    const tci = tc.index ?? 0;
                    if (tc.id) {
                        output.content.push({ type: "toolCall", id: tc.id, name: tc.function?.name ?? "", arguments: {} });
                        const ci = output.content.length - 1;
                        tcContentIdx.set(tci, ci);
                        tcJsonBufs.set(tci, "");
                        stream.push({ type: "toolcall_start", contentIndex: ci, partial: output });
                    }
                    if (tc.function?.arguments) {
                        const ci = tcContentIdx.get(tci);
                        if (ci === undefined)
                            continue;
                        const buf = (tcJsonBufs.get(tci) ?? "") + tc.function.arguments;
                        tcJsonBufs.set(tci, buf);
                        const block = output.content[ci];
                        if (block.type === "toolCall") {
                            try {
                                block.arguments = JSON.parse(buf);
                            }
                            catch { }
                        }
                        stream.push({ type: "toolcall_delta", contentIndex: ci, delta: tc.function.arguments, partial: output });
                    }
                }
            }
            if (choice.finish_reason === "stop")
                output.stopReason = "stop";
            else if (choice.finish_reason === "length")
                output.stopReason = "length";
            else if (choice.finish_reason === "tool_calls")
                output.stopReason = "toolUse";
        }
        // Finalize blocks
        for (let i = 0; i < output.content.length; i++) {
            if (output.content[i].type === "text")
                stream.push({ type: "text_end", contentIndex: i, content: output.content[i].text, partial: output });
        }
        for (const [tci, ci] of tcContentIdx) {
            const b = output.content[ci];
            if (b.type === "toolCall") {
                try {
                    b.arguments = JSON.parse(tcJsonBufs.get(tci) ?? "{}");
                }
                catch { }
                stream.push({ type: "toolcall_end", contentIndex: ci, toolCall: b, partial: output });
            }
        }
    })();
}
// =============================================================================
// Anthropic Messages API streaming
// =============================================================================
function streamAnthropic(model, context, options, output, stream, baseHost, auth) {
    return (async () => {
        const url = `${baseHost}/anthropic/v1/messages`;
        const body = {
            model: model.id,
            messages: toAnthropicMessages(context.messages),
            max_tokens: options?.maxTokens ?? model.maxTokens,
            stream: true,
        };
        if (context.systemPrompt)
            body.system = context.systemPrompt;
        if (context.tools?.length)
            body.tools = toAnthropicTools(context.tools);
        const token = await auth.getToken();
        // Anthropic route on Azure Foundry always uses "Authorization: Bearer"
        // regardless of auth type — api-key values are valid Bearer tokens here.
        const response = await fetch(url, {
            method: "POST",
            headers: {
                "Content-Type": "application/json",
                "Authorization": `Bearer ${token}`,
                "anthropic-version": "2023-06-01",
            },
            body: JSON.stringify(body),
            signal: options?.signal,
        });
        if (!response.ok) {
            const t = await response.text().catch(() => "");
            throw new Error(`Azure Foundry ${response.status}: ${t.slice(0, 500)}`);
        }
        if (!response.body)
            throw new Error("No response body");
        stream.push({ type: "start", partial: output });
        // Anthropic SSE events: message_start, content_block_start, content_block_delta, content_block_stop, message_delta, message_stop
        const blockIndices = new Map(); // anthropic block index → output.content index
        const tcJsonBufs = new Map();
        const reader = response.body.getReader();
        for await (const data of parseSSE(reader)) {
            let event;
            try {
                event = JSON.parse(data);
            }
            catch {
                continue;
            }
            if (event.type === "message_start" && event.message?.usage) {
                output.usage.input = event.message.usage.input_tokens ?? 0;
                output.usage.cacheRead = event.message.usage.cache_read_input_tokens ?? 0;
                output.usage.cacheWrite = event.message.usage.cache_creation_input_tokens ?? 0;
            }
            if (event.type === "content_block_start") {
                const cb = event.content_block;
                const anthropicIdx = event.index;
                if (cb.type === "text") {
                    output.content.push({ type: "text", text: "" });
                    const ci = output.content.length - 1;
                    blockIndices.set(anthropicIdx, ci);
                    stream.push({ type: "text_start", contentIndex: ci, partial: output });
                }
                else if (cb.type === "thinking") {
                    output.content.push({ type: "thinking", thinking: "", thinkingSignature: "" });
                    blockIndices.set(anthropicIdx, output.content.length - 1);
                    stream.push({ type: "thinking_start", contentIndex: output.content.length - 1, partial: output });
                }
                else if (cb.type === "tool_use") {
                    output.content.push({ type: "toolCall", id: cb.id, name: cb.name, arguments: {} });
                    const ci = output.content.length - 1;
                    blockIndices.set(anthropicIdx, ci);
                    tcJsonBufs.set(anthropicIdx, "");
                    stream.push({ type: "toolcall_start", contentIndex: ci, partial: output });
                }
            }
            if (event.type === "content_block_delta") {
                const ci = blockIndices.get(event.index);
                if (ci === undefined)
                    continue;
                const block = output.content[ci];
                const d = event.delta;
                if (d.type === "text_delta" && block.type === "text") {
                    block.text += d.text;
                    stream.push({ type: "text_delta", contentIndex: ci, delta: d.text, partial: output });
                }
                else if (d.type === "thinking_delta" && block.type === "thinking") {
                    block.thinking += d.thinking;
                    stream.push({ type: "thinking_delta", contentIndex: ci, delta: d.thinking, partial: output });
                }
                else if (d.type === "signature_delta" && block.type === "thinking") {
                    block.thinkingSignature = (block.thinkingSignature ?? "") + d.signature;
                }
                else if (d.type === "input_json_delta" && block.type === "toolCall") {
                    const buf = (tcJsonBufs.get(event.index) ?? "") + d.partial_json;
                    tcJsonBufs.set(event.index, buf);
                    try {
                        block.arguments = JSON.parse(buf);
                    }
                    catch { }
                    stream.push({ type: "toolcall_delta", contentIndex: ci, delta: d.partial_json, partial: output });
                }
            }
            if (event.type === "content_block_stop") {
                const ci = blockIndices.get(event.index);
                if (ci === undefined)
                    continue;
                const block = output.content[ci];
                if (block.type === "text")
                    stream.push({ type: "text_end", contentIndex: ci, content: block.text, partial: output });
                else if (block.type === "thinking")
                    stream.push({ type: "thinking_end", contentIndex: ci, content: block.thinking, partial: output });
                else if (block.type === "toolCall") {
                    try {
                        block.arguments = JSON.parse(tcJsonBufs.get(event.index) ?? "{}");
                    }
                    catch { }
                    stream.push({ type: "toolcall_end", contentIndex: ci, toolCall: block, partial: output });
                }
            }
            if (event.type === "message_delta") {
                if (event.usage) {
                    output.usage.output = event.usage.output_tokens ?? 0;
                    output.usage.totalTokens = output.usage.input + output.usage.output + output.usage.cacheRead + output.usage.cacheWrite;
                    calculateCost(model, output.usage);
                }
                const sr = event.delta?.stop_reason;
                if (sr === "end_turn" || sr === "stop_sequence")
                    output.stopReason = "stop";
                else if (sr === "max_tokens")
                    output.stopReason = "length";
                else if (sr === "tool_use")
                    output.stopReason = "toolUse";
            }
        }
    })();
}
// =============================================================================
// Unified streamSimple — routes based on publisher
// =============================================================================
function streamAzureFoundry(model, context, options) {
    const stream = createAssistantMessageEventStream();
    (async () => {
        const output = {
            role: "assistant",
            content: [],
            api: model.api,
            provider: model.provider,
            model: model.id,
            usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
            stopReason: "stop",
            timestamp: Date.now(),
        };
        try {
            const baseHost = new URL(model.baseUrl).origin;
            const publisher = publisherMap.get(model.id) ?? "";
            // Resolve auth: use registered provider auth, fall back to api-key from options.
            const auth = providerAuthMap.get(model.provider)
                ?? { type: "api-key", getToken: () => Promise.resolve(options?.apiKey ?? "") };
            if (publisher === "Anthropic") {
                await streamAnthropic(model, context, options, output, stream, baseHost, auth);
            }
            else {
                await streamOpenAI(model, context, options, output, stream, baseHost, auth);
            }
            stream.push({ type: "done", reason: output.stopReason, message: output });
            stream.end();
        }
        catch (error) {
            output.stopReason = options?.signal?.aborted ? "aborted" : "error";
            output.errorMessage = error instanceof Error ? error.message : String(error);
            stream.push({ type: "error", reason: output.stopReason, error: output });
            stream.end();
        }
    })();
    return stream;
}
// =============================================================================
// Extension Entry Point
// =============================================================================
export default async function (pi) {
    const config = loadConfig();
    const endpoint = `https://${config.resourceId}.services.ai.azure.com/api/projects/${config.projectId}`;
    // Discover deployments
    const url = `${endpoint}/deployments?api-version=v1`;
    console.log(`[Azure Foundry] Fetching deployments from: ${url}`);
    const getToken = makeTokenGetter(config.auth);
    console.log(`[Azure Foundry] Auth: ${config.auth.type}`);
    const token = await getToken();
    const response = await fetch(url, { headers: { Authorization: `Bearer ${token}` } });
    if (!response.ok) {
        const b = await response.text().catch(() => "");
        throw new Error(`Azure Foundry API ${response.status}: ${b.slice(0, 200)}`);
    }
    const data = (await response.json());
    const deployments = (data.value ?? []).filter((d) => d.capabilities?.chat_completion === "true");
    if (deployments.length === 0)
        throw new Error("No chat-capable deployments found");
    const summary = deployments.map((d) => `${d.name} (${d.modelPublisher})`).join(", ");
    console.log(`[Azure Foundry] Found ${deployments.length} deployment(s): ${summary}`);
    const providerId = "azure-foundry";
    // Store the auth context so streamAzureFoundry can build the right headers per-request.
    providerAuthMap.set(providerId, { type: config.auth.type, getToken });
    pi.registerProvider(providerId, {
        name: "Azure Foundry",
        baseUrl: endpoint,
        // For api-key auth, store the key so pi can display it in settings.
        // For azure-identity, leave empty — tokens are fetched at request time.
        apiKey: config.auth.type === "api-key" ? config.auth.apiKey : "",
        api: "azure-foundry",
        streamSimple: streamAzureFoundry,
        models: deployments.map(deploymentToModel),
    });
    console.log(`[Azure Foundry] ✓ Registered ${deployments.length} model(s)`);
}
//# sourceMappingURL=index.js.map