import {
  AutoTokenizer,
  AutoProcessor,
  AutoModelForCausalLM,
  AutoModelForImageTextToText,
  TextStreamer,
  RawImage,
  env,
} from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.8.1';

env.allowLocalModels = false;
env.useBrowserCache = true;

// Mapping of model prefixes to their corresponding model classes
const MODEL_CLASSES = {
    'onnx-community/granite-4.0-micro-ONNX-web': {
        model: AutoModelForCausalLM,
        processor: AutoTokenizer
    },
    'mistralai/Ministral-3-3B-Instruct-2512-ONNX': {
        model: AutoModelForImageTextToText,
        processor: AutoProcessor
    },
};

class TextGenerationPipeline {
  static processor = null; // Can be a Processor or a Tokenizer
  static model = null;
  static model_id = null; // Track loaded model
  static model_class = null; // Track loaded model class

  static async getInstance(model_id, progress_callback = null) {
    if (model_id && model_id !== this.model_id) {
        // Force reload if a new model is requested
        this.processor = null;
        this.model = null;
        this.model_class = null;
    }

    if (this.processor === null || this.model === null) {
      if (!model_id) throw new Error("Model not loaded and no model ID provided to load.");
      
      const modelInfo = MODEL_CLASSES[model_id] || { model: AutoModelForCausalLM, processor: AutoTokenizer };
      this.model_class = modelInfo.model;
      const ProcessorClass = modelInfo.processor;

      this.processor = await ProcessorClass.from_pretrained(model_id, {
        progress_callback,
      });

      // Set the image processor size for vision models
      if (this.model_class === AutoModelForImageTextToText && this.processor.image_processor) {
        this.processor.image_processor.size = { longest_edge: 480 };
      }

      const modelConfig = {
        device: "webgpu",
        progress_callback,
      };

      if (this.model_class === AutoModelForImageTextToText) {
        modelConfig.dtype = {
          embed_tokens: "fp16",
          vision_encoder: "q4",
          decoder_model_merged: "q4f16",
        };
      } else {
        modelConfig.dtype = "q4f16";
      }

      this.model = await this.model_class.from_pretrained(model_id, modelConfig);

      this.model_id = model_id;
    }
    return [this.processor, this.model, this.model_class];
  }
}


self.addEventListener("message", async (e) => {
  const { type, data } = e.data;

  switch (type) {
    case 'load':
      await loadModel(data.model_id);
      break;
    case 'generate':
      await generateText(data);
      break;
  }
});


async function loadModel(model_id) {
    self.postMessage({
        status: "loading",
        message: "Loading model...",
    });

    try {
        const [processor, model, model_class] = await TextGenerationPipeline.getInstance(model_id, (x) => {
            self.postMessage({ status: 'progress', data: x });
        });

        self.postMessage({
            status: "loading",
            message: "Compiling shaders and warming up model...",
        });
        
        // Perform a dry run to compile shaders
        if (model_class === AutoModelForImageTextToText) {
            const image = new RawImage(new Uint8ClampedArray(480 * 480 * 4).fill(0), 480, 480, 4);
            const messages = [
                { role: "system", content: "You are a helpful AI assistant." },
                { role: "user", content: "[IMG]a" }
            ];
            self.postMessage({ status: 'log', message: 'Warming up model... (1/3) Applying chat template' });
            let prompt, inputs;
            try {
                prompt = processor.apply_chat_template(messages);
            } catch (e) {
                self.postMessage({ status: 'log', message: 'ERROR during apply_chat_template' });
                throw e;
            }

            self.postMessage({ status: 'log', message: 'Warming up model... (2/3) Processing inputs' });
            try {
                inputs = await processor(image, prompt);
            } catch (e) {
                self.postMessage({ status: 'log', message: 'ERROR during processor(image, prompt)' });
                throw e;
            }

            self.postMessage({ status: 'log', message: 'Warming up model... (3/3) Generating' });
            try {
                await model.generate({ ...inputs, max_new_tokens: 1 });
            } catch (e) {
                self.postMessage({ status: 'log', message: 'ERROR during model.generate' });
                throw e;
            }
        } else {
            const inputs = processor("a");
            await model.generate({ ...inputs, max_new_tokens: 1 });
        }

        self.postMessage({ status: "ready", message: `Successfully loaded ${model_id}` });

    } catch (err) {
        handleError(err);
    }
}

async function generateText({ messages, image }) {
    try {
        const [processor, model, model_class] = await TextGenerationPipeline.getInstance();

        let inputs;
        if (model_class === AutoModelForImageTextToText) {
            
            let final_image;
            // Create a shallow copy of the messages to avoid mutation
            const processedMessages = JSON.parse(JSON.stringify(messages));

            if (image) {
                // Use the user-uploaded image
                final_image = new RawImage(image.data, image.width, image.height, 4);
                
                // Prepend [IMG] to the last user message
                const lastUserMessage = processedMessages.filter(x => x.role === 'user').pop();
                if (lastUserMessage) {
                    lastUserMessage.content = `[IMG]${lastUserMessage.content}`;
                }

            } else {
                // Fallback to a dummy black image for text-only queries
                final_image = new RawImage(new Uint8ClampedArray(480 * 480 * 4).fill(0), 480, 480, 4);
            }

            const prompt = processor.apply_chat_template(processedMessages, { add_generation_prompt: true });
            inputs = await processor(final_image, prompt);

        } else {
            // Standard text-only chat model
            inputs = processor.apply_chat_template(messages, {
                add_generation_prompt: true,
                return_dict: true,
            });
        }


        const tokenizer = model_class === AutoModelForImageTextToText ? processor.tokenizer : processor;
        const streamer = new TextStreamer(tokenizer, {
            skip_prompt: true,
            skip_special_tokens: true,
            callback_function: (text) => {
                self.postMessage({ status: 'stream-chunk', text: text });
            },
        });

        self.postMessage({ status: 'log', message: 'Generating response...' });

        const output = await model.generate({
            ...inputs,
            max_new_tokens: 512,
            temperature: 0.7,
            do_sample: true,
            top_k: 20,
            streamer: streamer,
            return_dict_in_generate: true,
        });
        
        const decoded = processor.batch_decode(output.sequences, { skip_special_tokens: true });
        
        self.postMessage({ status: 'complete', output: decoded[0] });

    } catch (err) {
        handleError(err);
    }
}

async function handleError(err) {
    // Stringify the error to ensure all information is sent
    const errorObject = {
        message: err.message || 'No error message',
        stack: err.stack || 'No stack trace',
        name: err.name || 'No error name',
        // Include any other properties that might be useful
        ...err
    };
    const errorMsg = JSON.stringify(errorObject, Object.getOwnPropertyNames(err));

    self.postMessage({ status: 'error', message: errorMsg, stack: err.stack });
}