import { BaseService, type BaseServiceContext } from '@learnitive/base-utils';
import { parseBibTeX } from '@learnitive/bibliography';
import { SERVICE_NAME } from '@learnitive/constants';
import {
  buildProjectMeaningChunks,
  createProjectMeaningReuseLookup,
  EMBEDDING_DIMENSIONS,
  EMBEDDING_MODEL_VERSION,
  PROJECT_MEANING_SCHEMA_VERSION,
  rankProjectMeaning,
  reusableProjectMeaningVectorIndex,
  type EmbeddingCapabilities,
  type EmbeddingEngine,
  type EmbeddingModelState,
  type EmbeddingPlatform,
  type ProjectMeaningSearchResult,
  type StoredProjectMeaningChunk,
  type StoredProjectMeaningIndex,
} from '@learnitive/project-meaning';
import {
  buildProjectSemanticIndex,
  type ProjectSemanticSnapshot,
} from '@learnitive/project-semantic-index';
import {
  projectFilePath,
  type ProjectFilePath,
} from '@learnitive/project-storage';
import { WsPath } from '@learnitive/ws-path';
import type { LatexSourceService } from './latex-source-service';
import type { ProjectStorageService } from './project-storage-service';
import type { WorkspaceOpsService } from './workspace-ops-service';

const INDEXABLE_EXTENSIONS = /\.(?:tex|ltx|sty|cls)$/i;
const BIBLIOGRAPHY_EXTENSION = /\.bib$/i;
const decoder = new TextDecoder('utf-8', { fatal: true });
const encoder = new TextEncoder();
const MEANING_METADATA_KEY = 'learnitiveProjectMeaning';

type WorkspaceIndexState = {
  snapshot?: ProjectSemanticSnapshot;
  loading?: Promise<ProjectSemanticSnapshot>;
  refreshRequested?: boolean;
  listeners: Set<() => void>;
  refreshTimer?: ReturnType<typeof setTimeout>;
};

type Config = {
  createEmbeddingPlatform: () => EmbeddingPlatform;
};

export type ProjectMeaningStatus = {
  capabilities: EmbeddingCapabilities;
  model: EmbeddingModelState;
  index?: {
    passages: number;
    updatedAt: number;
    dimensions: number;
    modelVersion: string;
  };
};

export type ProjectMeaningIndexProgress = {
  phase: 'preparing' | 'hashing' | 'embedding' | 'saving' | 'ready';
  completed: number;
  total: number;
  reused: number;
};

function isMeaningMetadata(value: unknown): value is { version: 1; id: string } {
  return Boolean(
    value &&
    typeof value === 'object' &&
    (value as { version?: unknown }).version === 1 &&
    typeof (value as { id?: unknown }).id === 'string' &&
    (value as { id: string }).id.length > 8,
  );
}

async function sha256(value: string): Promise<string> {
  const digest = await crypto.subtle.digest('SHA-256', encoder.encode(value));
  return [...new Uint8Array(digest)].map((byte) => byte.toString(16).padStart(2, '0')).join('');
}

function vectorFrom(index: StoredProjectMeaningIndex, position: number): Float32Array {
  const start = position * index.dimensions * Float32Array.BYTES_PER_ELEMENT;
  return new Float32Array(index.vectors.slice(start, start + index.dimensions * Float32Array.BYTES_PER_ELEMENT));
}

function lineAtOffset(source: string, offset: number): number {
  let line = 1;
  const limit = Math.max(0, Math.min(offset, source.length));
  for (let index = 0; index < limit; index += 1) {
    if (source.charCodeAt(index) === 10) line += 1;
  }
  return line;
}

/** Shared project-wide LaTeX meaning used by Outline, Visual, Doctor and AI. */
export class ProjectSemanticService extends BaseService {
  static deps = ['projectStorage', 'latexSource', 'workspaceOps'] as const;
  private readonly workspaces = new Map<string, WorkspaceIndexState>();
  private readonly embedding: EmbeddingPlatform;
  private embeddingEngine?: EmbeddingEngine;

  constructor(
    context: BaseServiceContext,
    private readonly dependencies: {
      projectStorage: ProjectStorageService;
      latexSource: LatexSourceService;
      workspaceOps: WorkspaceOpsService;
    },
    private readonly config: Config,
  ) {
    super(SERVICE_NAME.projectSemanticService, context, dependencies);
    this.embedding = config.createEmbeddingPlatform();
  }

  async hookMount(): Promise<void> {
    const unsubscribe = this.dependencies.latexSource.subscribeToSaveStatus(() => {
      for (const wsName of this.workspaces.keys()) this.scheduleRefresh(wsName);
    });
    this.addCleanup(unsubscribe);
    this.addCleanup(() => {
      for (const state of this.workspaces.values()) {
        if (state.refreshTimer) clearTimeout(state.refreshTimer);
      }
      this.workspaces.clear();
      void this.embeddingEngine?.unload();
      this.embeddingEngine = undefined;
      this.embedding.dispose();
    });
  }

  private state(wsName: string): WorkspaceIndexState {
    let state = this.workspaces.get(wsName);
    if (!state) {
      state = { listeners: new Set() };
      this.workspaces.set(wsName, state);
    }
    return state;
  }

  getSnapshot(wsName: string): ProjectSemanticSnapshot | undefined {
    return this.workspaces.get(wsName)?.snapshot;
  }

  subscribe(wsName: string, listener: () => void): () => void {
    const state = this.state(wsName);
    state.listeners.add(listener);
    return () => state.listeners.delete(listener);
  }

  async refresh(wsName: string, signal?: AbortSignal): Promise<ProjectSemanticSnapshot> {
    await this.mountPromise;
    signal?.throwIfAborted();
    const state = this.state(wsName);
    if (state.loading) {
      state.refreshRequested = true;
      return state.loading;
    }
    state.loading = (async () => {
      let snapshot: ProjectSemanticSnapshot;
      do {
        state.refreshRequested = false;
        snapshot = await this.build(wsName, signal);
      } while (state.refreshRequested && !signal?.aborted);
      return snapshot;
    })();
    try {
      const snapshot = await state.loading;
      state.snapshot = snapshot;
      for (const listener of state.listeners) listener();
      return snapshot;
    } finally {
      state.loading = undefined;
    }
  }

  private scheduleRefresh(wsName: string): void {
    const state = this.state(wsName);
    if (state.refreshTimer) clearTimeout(state.refreshTimer);
    state.refreshTimer = setTimeout(() => {
      state.refreshTimer = undefined;
      void this.refresh(wsName).catch((error) => {
        console.warn('Unable to refresh Learnitive semantic index', error);
      });
    }, 2_000);
  }

  private async meaningId(wsName: string): Promise<string> {
    const metadata = await this.dependencies.workspaceOps.getWorkspaceMetadata(wsName);
    const existing = metadata[MEANING_METADATA_KEY];
    if (isMeaningMetadata(existing)) return existing.id;
    const id = crypto.randomUUID();
    await this.dependencies.workspaceOps.updateWorkspaceMetadata(wsName, (current) => ({
      ...current,
      [MEANING_METADATA_KEY]: { version: 1, id },
    }));
    return id;
  }

  async meaningStatus(wsName: string): Promise<ProjectMeaningStatus> {
    const [capabilities, model, memoryId] = await Promise.all([
      this.embedding.capabilities(),
      this.embedding.modelState(),
      this.meaningId(wsName),
    ]);
    const index = await this.embedding.store.load(memoryId);
    return {
      capabilities,
      model,
      ...(index &&
        index.schemaVersion === PROJECT_MEANING_SCHEMA_VERSION &&
        index.modelId === this.embedding.modelId &&
        index.modelVersion === EMBEDDING_MODEL_VERSION &&
        index.dimensions === EMBEDDING_DIMENSIONS
        ? {
            index: {
              passages: index.chunks.length,
              updatedAt: index.updatedAt,
              dimensions: index.dimensions,
              modelVersion: index.modelVersion,
            },
          }
        : {}),
    };
  }

  async prepareMeaningModel(
    onProgress?: (state: EmbeddingModelState) => void,
    signal?: AbortSignal,
  ): Promise<EmbeddingModelState> {
    return this.embedding.prepareModel(onProgress, signal);
  }

  async removeMeaningModel(): Promise<void> {
    await this.embeddingEngine?.unload().catch(() => undefined);
    this.embeddingEngine = undefined;
    await this.embedding.removeModel();
  }

  async clearMeaningIndex(wsName: string): Promise<void> {
    await this.embedding.store.clear(await this.meaningId(wsName));
  }

  private async engine(signal?: AbortSignal): Promise<EmbeddingEngine> {
    if (!this.embeddingEngine) this.embeddingEngine = this.embedding.createEngine();
    await this.embeddingEngine.load(signal);
    return this.embeddingEngine;
  }

  async indexMeaning(
    wsName: string,
    options: {
      signal?: AbortSignal;
      onProgress?: (progress: ProjectMeaningIndexProgress) => void;
    } = {},
  ): Promise<StoredProjectMeaningIndex> {
    await this.mountPromise;
    const { signal, onProgress } = options;
    signal?.throwIfAborted();
    const model = await this.embedding.modelState();
    if (model.status !== 'installed' && model.status !== 'ready') {
      throw new Error('Download EmbeddingGemma before indexing project meaning');
    }

    onProgress?.({ phase: 'preparing', completed: 0, total: 0, reused: 0 });
    const snapshot = await this.refresh(wsName, signal);
    const chunks = buildProjectMeaningChunks(snapshot);
    // Meaning results preserve exact source ranges. Recompute display lines
    // from canonical/open source so ordinary prose does not collapse to line 1
    // merely because it has no semantic command node of its own.
    const provider = await this.dependencies.projectStorage.getProvider(wsName);
    const sourceByFile = new Map<ProjectFilePath, string>();
    for (const file of new Set(chunks.map((chunk) => chunk.file))) {
      const wsPath = WsPath.fromParts(wsName, file).wsPath;
      const open = this.dependencies.latexSource.getSnapshot(wsPath);
      if (open) {
        sourceByFile.set(file, open.text);
        continue;
      }
      const data = await provider.readFile(file, signal);
      if (!data) continue;
      try {
        sourceByFile.set(file, decoder.decode(data));
      } catch {
        // The deterministic index already excluded invalid UTF-8 source.
      }
    }
    for (const chunk of chunks) {
      const source = sourceByFile.get(chunk.file);
      if (source) chunk.line = lineAtOffset(source, chunk.range.from);
    }
    const memoryId = await this.meaningId(wsName);
    const previous = await this.embedding.store.load(memoryId);
    const previousCompatible = previous &&
      previous.schemaVersion === PROJECT_MEANING_SCHEMA_VERSION &&
      previous.modelId === this.embedding.modelId &&
      previous.modelVersion === EMBEDDING_MODEL_VERSION &&
      previous.dimensions === EMBEDDING_DIMENSIONS;
    const reuseLookup = previousCompatible ? createProjectMeaningReuseLookup(previous) : undefined;

    const hashed: StoredProjectMeaningChunk[] = [];
    for (let i = 0; i < chunks.length; i += 1) {
      signal?.throwIfAborted();
      const chunk = chunks[i]!;
      const contentHash = await sha256(`${chunk.kind}\n${chunk.title ?? ''}\n${chunk.text}`);
      hashed.push({ ...chunk, contentHash });
      onProgress?.({ phase: 'hashing', completed: i + 1, total: chunks.length, reused: 0 });
    }

    const vectors = new Array<Float32Array | undefined>(hashed.length);
    const pending: Array<{ index: number; chunk: StoredProjectMeaningChunk }> = [];
    let reused = 0;
    for (let i = 0; i < hashed.length; i += 1) {
      const chunk = hashed[i]!;
      const reusableIndex = reuseLookup
        ? reusableProjectMeaningVectorIndex(reuseLookup, chunk)
        : undefined;
      if (reusableIndex !== undefined && previous) {
        vectors[i] = vectorFrom(previous, reusableIndex);
        reused += 1;
      } else {
        pending.push({ index: i, chunk });
      }
    }

    const embeddingEngine = pending.length ? await this.engine(signal) : undefined;
    const batchSize = 12;
    let embedded = 0;
    for (let offset = 0; offset < pending.length; offset += batchSize) {
      signal?.throwIfAborted();
      const batch = pending.slice(offset, offset + batchSize);
      const values = await embeddingEngine!.embedDocuments(batch.map(({ chunk }) => chunk.title ? `${chunk.title}\n${chunk.text}` : chunk.text), signal);
      if (values.length !== batch.length) throw new Error('Embedding engine returned an unexpected result count');
      values.forEach((vector, position) => {
        if (vector.length !== EMBEDDING_DIMENSIONS) throw new Error(`Unexpected EmbeddingGemma dimensions: ${vector.length}`);
        vectors[batch[position]!.index] = vector;
      });
      embedded += batch.length;
      onProgress?.({ phase: 'embedding', completed: embedded + reused, total: hashed.length, reused });
    }

    const contiguous = new Float32Array(hashed.length * EMBEDDING_DIMENSIONS);
    vectors.forEach((vector, index) => {
      if (!vector) throw new Error('Project meaning index is missing a vector');
      contiguous.set(vector, index * EMBEDDING_DIMENSIONS);
    });

    const next: StoredProjectMeaningIndex = {
      schemaVersion: PROJECT_MEANING_SCHEMA_VERSION,
      generationId: crypto.randomUUID(),
      modelId: this.embedding.modelId,
      modelVersion: EMBEDDING_MODEL_VERSION,
      dimensions: EMBEDDING_DIMENSIONS,
      updatedAt: Date.now(),
      chunks: hashed,
      vectors: contiguous.buffer,
    };
    onProgress?.({ phase: 'saving', completed: hashed.length, total: hashed.length, reused });
    await this.embedding.store.save(memoryId, next);
    onProgress?.({ phase: 'ready', completed: hashed.length, total: hashed.length, reused });
    return next;
  }

  async searchMeaning(
    wsName: string,
    query: string,
    options: { limit?: number; signal?: AbortSignal; onProgress?: (progress: ProjectMeaningIndexProgress) => void } = {},
  ): Promise<ProjectMeaningSearchResult[]> {
    const normalized = query.trim();
    if (!normalized) return [];
    const index = await this.indexMeaning(wsName, options);
    if (index.chunks.length === 0) return [];
    const engine = await this.engine(options.signal);
    const queryVector = await engine.embedQuery(normalized, options.signal);
    return rankProjectMeaning(index, queryVector, options.limit ?? 30);
  }

  private async build(wsName: string, signal?: AbortSignal): Promise<ProjectSemanticSnapshot> {
    const provider = await this.dependencies.projectStorage.getProvider(wsName);
    const config = await this.dependencies.projectStorage.getProjectConfig(wsName);
    const entries = await provider.listEntries(signal);
    const allFiles = new Set<ProjectFilePath>();
    const sources = new Map<ProjectFilePath, string>();
    const bibliographies = new Map<ProjectFilePath, ReturnType<typeof parseBibTeX>>();

    for (const entry of entries) {
      signal?.throwIfAborted();
      if (entry.kind !== 'file') continue;
      allFiles.add(entry.path);
      if (!INDEXABLE_EXTENSIONS.test(entry.path) && !BIBLIOGRAPHY_EXTENSION.test(entry.path)) continue;

      const wsPath = WsPath.fromParts(wsName, entry.path).wsPath;
      const open = this.dependencies.latexSource.getSnapshot(wsPath);
      let text: string | undefined;
      if (open) {
        text = open.text;
      } else {
        const data = await provider.readFile(entry.path, signal);
        if (!data) continue;
        try {
          text = decoder.decode(data);
        } catch {
          // A text-like file with non-UTF8 contents remains byte-safe in
          // storage; it is simply omitted from the semantic index.
        }
      }
      if (text === undefined) continue;
      if (BIBLIOGRAPHY_EXTENSION.test(entry.path)) {
        bibliographies.set(entry.path, parseBibTeX(text));
      } else {
        sources.set(entry.path, text);
      }
    }

    return buildProjectSemanticIndex(
      sources,
      projectFilePath(config.rootFile),
      allFiles,
      bibliographies,
    );
  }
}