import {
  LOCAL_AI_MODELS,
  type GenerationRequest,
  type LocalAiModelId,
  type LocalGenerationEngine,
  type ProjectAiProviderMessage,
} from '@learnitive/project-ai';

// Google documents LiteRT-LM's Web API as a normal browser ESM import. Do not
// wrap it in a module Worker: the preview runtime still uses importScripts()
// internally on some code paths, which is illegal inside a module Worker.
const LITERT_LM_MODULE_URL = 'https://cdn.jsdelivr.net/npm/@litert-lm/core@0.15.0/+esm';
const STORE_DIRECTORY = 'interactive-local-ai-v1';
const MODEL_DIRECTORY = 'models';

type LiteRtMessage = {
  role: string;
  content?: string | Array<{ type?: string; text?: string }>;
};

type LiteRtConversation = {
  sendMessageStreaming(message: LiteRtMessage | string): ReadableStream<LiteRtMessage>;
  cancel(): void;
  delete(): Promise<void>;
};

type LiteRtEngine = {
  createConversation(config?: { preface?: { messages?: LiteRtMessage[] } }): Promise<LiteRtConversation>;
  delete(): Promise<void>;
};

type LiteRtModule = {
  Engine: {
    create(settings: {
      model: ReadableStream<Uint8Array> | Blob | string;
      mainExecutorSettings?: { maxNumTokens?: number };
    }): Promise<LiteRtEngine>;
  };
};

let runtimeModule: Promise<LiteRtModule> | undefined;
let fetchReceiverGuardInstalled = false;

/**
 * LiteRT-LM currently reaches the browser fetch implementation through code
 * paths that may cache/extract `window.fetch` before invoking it. Chromium
 * treats the native Window method as receiver-sensitive and throws
 * "Illegal invocation" when such an extracted function is called. Install a
 * tiny stable wrapper once, before the LiteRT module is loaded, so both direct
 * and extracted calls preserve the Window/global receiver. This does not alter
 * request contents, routing, credentials, or Learnitive privacy policy.
 */
function ensureFetchReceiverSafe(): void {
  if (fetchReceiverGuardInstalled || typeof globalThis.fetch !== 'function') return;
  const nativeFetch = globalThis.fetch;
  const receiver = typeof window === 'undefined' ? globalThis : window;
  // Use a genuinely bound function, not a method-shaped wrapper. LiteRT-LM
  // and its transitive runtime may cache/extract `fetch`; a bound function
  // remains valid regardless of the receiver used by the caller.
  globalThis.fetch = nativeFetch.bind(receiver) as typeof globalThis.fetch;
  fetchReceiverGuardInstalled = true;
}

function descriptor(modelId: LocalAiModelId) {
  const model = LOCAL_AI_MODELS.find((item) => item.id === modelId);
  if (!model) throw new Error('Unknown local AI model');
  return model;
}

function storageManager(): StorageManager & { getDirectory?: () => Promise<FileSystemDirectoryHandle> } {
  return navigator.storage as StorageManager & { getDirectory?: () => Promise<FileSystemDirectoryHandle> };
}

async function modelFile(modelId: LocalAiModelId): Promise<File> {
  const getDirectory = storageManager().getDirectory;
  if (!getDirectory) throw new Error('Origin-private model storage is unavailable');
  const root = await getDirectory.call(navigator.storage);
  const store = await root.getDirectoryHandle(STORE_DIRECTORY, { create: false });
  const models = await store.getDirectoryHandle(MODEL_DIRECTORY, { create: false });
  const handle = await models.getFileHandle(`${modelId}.litertlm`);
  const file = await handle.getFile();
  if (file.size <= 0) throw new Error('The installed local Gemma model is empty; remove and download it again');
  return file;
}

async function runtime(): Promise<LiteRtModule> {
  ensureFetchReceiverSafe();
  runtimeModule ??= (import(/* @vite-ignore */ LITERT_LM_MODULE_URL) as Promise<LiteRtModule>).catch((error) => {
    runtimeModule = undefined;
    throw error;
  });
  return runtimeModule;
}

function toLiteRtMessage(message: ProjectAiProviderMessage): LiteRtMessage {
  return { role: message.role, content: message.content };
}


function configuredContextWindowTokens(modelId: LocalAiModelId): number {
  const deviceMemory = Number((navigator as Navigator & { deviceMemory?: number }).deviceMemory ?? 0);
  return modelId === 'gemma-4-e2b-it-web'
    ? (deviceMemory > 0 && deviceMemory <= 4 ? 1024 : 2048)
    : 3072;
}

function messageText(message: LiteRtMessage): string {
  if (typeof message.content === 'string') return message.content;
  if (!Array.isArray(message.content)) return '';
  return message.content
    .filter((part) => part?.type === 'text' && typeof part.text === 'string')
    .map((part) => part.text ?? '')
    .join('');
}

class BrowserLiteRtGenerationEngine implements LocalGenerationEngine {
  private engine: LiteRtEngine | undefined;
  private activeConversation: LiteRtConversation | undefined;

  constructor(private readonly modelId: LocalAiModelId) {}

  contextWindowTokens(): number {
    return configuredContextWindowTokens(this.modelId);
  }

  async load(signal?: AbortSignal): Promise<void> {
    if (this.engine) return;
    ensureFetchReceiverSafe();
    signal?.throwIfAborted();
    // Validate the configured model before paying the runtime startup cost.
    descriptor(this.modelId);
    if (!('gpu' in navigator)) throw new Error('WebGPU is unavailable in this browser');
    const file = await modelFile(this.modelId);
    signal?.throwIfAborted();
    const module = await runtime();
    signal?.throwIfAborted();
    try {
      const deviceMemory = Number((navigator as Navigator & { deviceMemory?: number }).deviceMemory ?? 0);
      if (this.modelId === 'gemma-4-e4b-it-web' && deviceMemory > 0 && deviceMemory <= 8) {
        throw new Error('Gemma 4 E4B is disabled on browsers reporting 8 GB RAM or less because loading it can exhaust the tab/GPU process. Use Gemma 4 E2B instead.');
      }
      // Keep the KV/cache budget deliberately conservative. LiteRT-LM and the
      // model already consume substantial GPU/system memory; an 8k context can
      // make Chromium kill the renderer instead of throwing a recoverable error.
      const maxNumTokens = configuredContextWindowTokens(this.modelId);
      this.engine = await module.Engine.create({
        model: file.stream(),
        mainExecutorSettings: { maxNumTokens },
      });
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      if (/importScripts.*module scripts/iu.test(message)) {
        throw new Error(
          'LiteRT-LM was started in an incompatible worker context. Reload Learnitive after installing this build; inference now runs from the supported browser-module context.',
          { cause: error },
        );
      }
      throw error;
    }
  }

  async unload(): Promise<void> {
    this.activeConversation?.cancel();
    if (this.activeConversation) await this.activeConversation.delete().catch(() => undefined);
    this.activeConversation = undefined;
    if (this.engine) await this.engine.delete().catch(() => undefined);
    this.engine = undefined;
  }

  async *generate(request: GenerationRequest, signal?: AbortSignal) {
    await this.load(signal);
    if (!this.engine) throw new Error('Local AI engine did not initialize');
    if (this.activeConversation) throw new Error('Local AI is already generating a response');
    const messages = request.messages.map(toLiteRtMessage);
    if (messages.length === 0) throw new Error('Local AI request has no messages');
    const final = messages.at(-1);
    if (!final) throw new Error('Local AI request has no final message');
    const preface = messages.slice(0, -1);
    const conversation = await this.engine.createConversation({
      ...(preface.length ? { preface: { messages: preface } } : {}),
    });
    const maxOutputChars = request.maxTokens
      ? Math.max(192, Math.trunc(request.maxTokens) * 3)
      : undefined;
    let emittedChars = 0;
    this.activeConversation = conversation;
    const onAbort = () => conversation.cancel();
    signal?.addEventListener('abort', onAbort, { once: true });
    try {
      const stream = conversation.sendMessageStreaming(final);
      const reader = stream.getReader();
      try {
        while (true) {
          signal?.throwIfAborted();
          const { value, done } = await reader.read();
          if (done) break;
          const text = value ? messageText(value) : '';
          if (!text) continue;
          if (maxOutputChars === undefined) {
            yield { text };
            continue;
          }
          const remaining = maxOutputChars - emittedChars;
          if (remaining <= 0) {
            conversation.cancel();
            break;
          }
          const boundedText = text.length <= remaining ? text : text.slice(0, remaining);
          if (boundedText) {
            emittedChars += boundedText.length;
            yield { text: boundedText };
          }
          if (boundedText.length < text.length || emittedChars >= maxOutputChars) {
            conversation.cancel();
            break;
          }
        }
      } finally {
        reader.releaseLock();
      }
    } finally {
      signal?.removeEventListener('abort', onAbort);
      if (this.activeConversation === conversation) this.activeConversation = undefined;
      await conversation.delete().catch(() => undefined);
      const deviceMemory = Number((navigator as Navigator & { deviceMemory?: number }).deviceMemory ?? 0);
      if (deviceMemory > 0 && deviceMemory <= 4 && this.engine) {
        // On low-memory devices, retaining the model after a response is more
        // likely to trigger a browser/GPU-process crash than to save useful time.
        await this.engine.delete().catch(() => undefined);
        this.engine = undefined;
      }
    }
  }
}

export function createBrowserLiteRtGenerationEngine(modelId: LocalAiModelId): LocalGenerationEngine {
  return new BrowserLiteRtGenerationEngine(modelId);
}