import { ANSWER_SCHEMAS } from './auditor-core.js';

// Loads the auditor model in WebLLM (runs on the device's GPU through WebGPU), for the desktop app and the phone page.
// The model is downloaded once and cached by the browser, so later runs work offline.

// The fine-tuned Devseis auditor model (https://huggingface.co/Devseis/endpoint-auditor-0.5b), published per version
// as WebLLM/MLC weights. It has the same architecture as Qwen2.5-0.5B-Instruct, so WebLLM's prebuilt runtime library
// for that model is reused and only the weights differ. Set MODEL_VERSION to null to use the base model.
export const MODEL_VERSION = 'v0.3';
const MODEL_REPO = (quant, version) => `https://huggingface.co/Devseis/endpoint-auditor-0.5b-${quant}-MLC/resolve/${version}/`;
// Builds: q0f16 = unquantised (about 1 GB, best quality; desktop), q4f16_1 = 4-bit (about 280 MB; phones),
// q4f32_1 = 4-bit for GPUs without 16-bit maths (shader-f16).
const BASE_MODEL = quant => `Qwen2.5-0.5B-Instruct-${quant}-MLC`;
// Output length per task (the AI-tool classification is the longest answer).
const MAX_TOKENS = { finding: 900, summary: 1200, ai_classification: 1600 };

// webllm: the @mlc-ai/web-llm module (bundled in the app, from a CDN on the phone page).
// precision: 'full' (desktop: q0f16 when the GPU supports 16-bit maths) or 'compact' (phones: 4-bit).
// version: model tag to load (from app/model-version.json on the Space), or { version, version_4bit } to use a different
// tag for the 4-bit builds; defaults to MODEL_VERSION.
export async function loadModelWith(webllm, onProgress, precision = 'compact', version = MODEL_VERSION) {
  const { CreateMLCEngine, prebuiltAppConfig } = webllm;
  if (!('gpu' in navigator)) throw new Error('WebGPU is not available on this device.');
  const adapter = await navigator.gpu.requestAdapter();
  if (!adapter) throw new Error('No WebGPU-capable graphics adapter was found.');

  const f16 = adapter.features.has('shader-f16');
  const quant = !f16 ? 'q4f32_1' : precision === 'full' ? 'q0f16' : 'q4f16_1';
  if (version && typeof version === 'object') version = (quant.startsWith('q4') ? version.version_4bit : null) || version.version || MODEL_VERSION;
  const base = prebuiltAppConfig.model_list.find(m => m.model_id === BASE_MODEL(quant));
  if (!base) throw new Error('The base model is not in this WebLLM version.');
  const custom = version ? { ...base, model: MODEL_REPO(quant, version), model_id: `devseis-auditor-0.5b-${version}-${quant}` } : null;
  const appConfig = { ...prebuiltAppConfig, model_list: custom ? [...prebuiltAppConfig.model_list, custom] : prebuiltAppConfig.model_list };
  const initProgressCallback = report => onProgress({ progress: report.progress, text: report.text });

  let engine, modelId, fineTuned = false;
  if (custom) {
    try {
      engine = await CreateMLCEngine(custom.model_id, { appConfig, initProgressCallback });
      modelId = custom.model_id;
      fineTuned = true;
    } catch (error) {
      console.warn(`Devseis model ${version} could not be loaded (${error.message}); using the base model.`);
    }
  }
  if (!engine) {
    engine = await CreateMLCEngine(base.model_id, { appConfig, initProgressCallback });
    modelId = base.model_id;
  }
  return {
    name: modelId,
    fineTuned,
    // Constrained decoding: the reply is always valid JSON in the answer schema of the task.
    async generate(system, prompt, task) {
      const reply = await engine.chat.completions.create({
        messages: [{ role: 'system', content: system }, { role: 'user', content: prompt }],
        temperature: 0,
        max_tokens: MAX_TOKENS[task] || 900,
        response_format: { type: 'json_object', schema: JSON.stringify(ANSWER_SCHEMAS[task]) },
      });
      return reply.choices[0].message.content;
    },
  };
}