// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

// Shared NuGet download and extraction utilities for install scripts.

'use strict';

const fs = require('fs');
const path = require('path');
const os = require('os');
const https = require('https');
const AdmZip = require('adm-zip');

const PLATFORM_MAP = {
  'win32-x64': 'win-x64',
  'win32-arm64': 'win-arm64',
  'linux-x64': 'linux-x64',
  'darwin-arm64': 'osx-arm64',
};
const platformKey = `${os.platform()}-${os.arch()}`;
const RID = PLATFORM_MAP[platformKey];
// Install binaries into foundry-local-core/<platform> inside the package root.
const BIN_DIR = path.join(__dirname, '..', 'foundry-local-core', platformKey);
const EXT = os.platform() === 'win32' ? '.dll' : os.platform() === 'darwin' ? '.dylib' : '.so';

const REQUIRED_FILES = [
  `Microsoft.AI.Foundry.Local.Core${EXT}`,
  `${os.platform() === 'win32' ? '' : 'lib'}onnxruntime${EXT}`,
  `${os.platform() === 'win32' ? '' : 'lib'}onnxruntime-genai${EXT}`,
];

const NUGET_FEED = 'https://api.nuget.org/v3/index.json';
const ORT_NIGHTLY_FEED = 'https://pkgs.dev.azure.com/aiinfra/PublicPackages/_packaging/ORT-Nightly/nuget/v3/index.json';

// --- Download helpers ---

async function downloadWithRetryAndRedirects(url, destStream = null) {
    const maxRedirects = 5;
    let currentUrl = url;
    let redirects = 0;

    while (redirects < maxRedirects) {
        const response = await new Promise((resolve, reject) => {
            https.get(currentUrl, (res) => resolve(res))
                 .on('error', reject);
        });

        if (response.statusCode >= 300 && response.statusCode < 400 && response.headers.location) {
            currentUrl = response.headers.location;
            response.resume();
            redirects++;
            console.log(`  Following redirect to ${new URL(currentUrl).host}...`);
            continue;
        }

        if (response.statusCode !== 200) {
            throw new Error(`Download failed with status ${response.statusCode}: ${currentUrl}`);
        }

        if (destStream) {
            response.pipe(destStream);
            return new Promise((resolve, reject) => {
                destStream.on('finish', resolve);
                destStream.on('error', reject);
                response.on('error', reject);
            });
        } else {
            let data = '';
            response.on('data', chunk => data += chunk);
            return new Promise((resolve, reject) => {
                response.on('end', () => resolve(data));
                response.on('error', reject);
            });
        }
    }
    throw new Error('Too many redirects');
}

async function downloadJson(url) {
    return JSON.parse(await downloadWithRetryAndRedirects(url));
}

async function downloadFile(url, dest) {
    const file = fs.createWriteStream(dest);
    try {
        await downloadWithRetryAndRedirects(url, file);
        file.close();
    } catch (e) {
        file.close();
        if (fs.existsSync(dest)) fs.unlinkSync(dest);
        throw e;
    }
}

const serviceIndexCache = new Map();

async function getBaseAddress(feedUrl) {
    if (!serviceIndexCache.has(feedUrl)) {
        serviceIndexCache.set(feedUrl, await downloadJson(feedUrl));
    }
    const resources = serviceIndexCache.get(feedUrl).resources || [];
    const res = resources.find(r => r['@type'] && r['@type'].startsWith('PackageBaseAddress/3.0.0'));
    if (!res) throw new Error('Could not find PackageBaseAddress/3.0.0 in NuGet feed.');
    const baseAddress = res['@id'];
    return baseAddress.endsWith('/') ? baseAddress : baseAddress + '/';
}

async function installPackage(artifact, tempDir, binDir, skipIfPresent) {
    const pkgName = artifact.name;
    const pkgVer = artifact.version;

    // Skip download if this package's main native binary is already present
    // (e.g. pre-populated by CI from a locally-built artifact).
    // Callers pass skipIfPresent=false when overriding (e.g. WinML over standard).
    if (skipIfPresent) {
        const prefix = os.platform() === 'win32' ? '' : 'lib';
        let expectedFile;
        if (pkgName.includes('Foundry.Local.Core')) {
            expectedFile = `Microsoft.AI.Foundry.Local.Core${EXT}`;
        } else if (pkgName.includes('OnnxRuntimeGenAI')) {
            expectedFile = `${prefix}onnxruntime-genai${EXT}`;
        } else if (pkgName.includes('OnnxRuntime')) {
            expectedFile = `${prefix}onnxruntime${EXT}`;
        }
        if (expectedFile && fs.existsSync(path.join(binDir, expectedFile))) {
            console.log(`  ${pkgName}: already present, skipping download.`);
            return;
        }
    }

    const baseAddress = await getBaseAddress(artifact.feed);
    const nameLower = pkgName.toLowerCase();
    const verLower = pkgVer.toLowerCase();
    const downloadUrl = `${baseAddress}${nameLower}/${verLower}/${nameLower}.${verLower}.nupkg`;

    const nupkgPath = path.join(tempDir, `${pkgName}.${pkgVer}.nupkg`);
    console.log(`  Downloading ${pkgName} ${pkgVer}...`);
    await downloadFile(downloadUrl, nupkgPath);

    console.log(`  Extracting...`);
    const zip = new AdmZip(nupkgPath);
    const targetPathPrefix = `runtimes/${RID}/native/`.toLowerCase();
    const entries = zip.getEntries().filter(e => {
        const p = e.entryName.toLowerCase();
        return p.includes(targetPathPrefix) && p.endsWith(EXT);
    });

    if (entries.length > 0) {
        entries.forEach(entry => {
            zip.extractEntryTo(entry, binDir, false, true);
            console.log(`    Extracted ${entry.name}`);
        });
    } else {
        console.warn(`    No files found for RID ${RID} in ${pkgName}.`);
    }

    // Write a metadata package.json with version info for diagnostics
    if (pkgName.startsWith('Microsoft.AI.Foundry.Local.Core')) {
        const pkgJsonPath = path.join(binDir, 'package.json');
        const pkgContent = {
            name: `@foundry-local-core/${platformKey}`,
            version: pkgVer,
            description: `Native binaries for Foundry Local SDK (${platformKey})`,
            private: true
        };
        fs.writeFileSync(pkgJsonPath, JSON.stringify(pkgContent, null, 2));
    }
}

async function runInstall(artifacts, options) {
    if (!RID) {
        console.warn(`[foundry-local] Unsupported platform: ${platformKey}. Skipping.`);
        return;
    }

    const binDir = (options && options.binDir) || BIN_DIR;
    // When a custom binDir is provided (e.g. WinML overriding standard),
    // don't skip packages whose output files already exist — we need to
    // overwrite them with the variant's binaries.
    const skipIfPresent = !(options && options.binDir);

    console.log(`[foundry-local] Installing native libraries for ${RID}...`);
    fs.mkdirSync(binDir, { recursive: true });

    const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'foundry-install-'));
    try {
        for (const artifact of artifacts) {
            await installPackage(artifact, tempDir, binDir, skipIfPresent);
        }
        console.log('[foundry-local] Installation complete.');
    } finally {
        try { fs.rmSync(tempDir, { recursive: true, force: true }); } catch {}
    }
}

module.exports = { NUGET_FEED, ORT_NIGHTLY_FEED, runInstall };
