import type { ProjectSemanticSnapshot, SourceRange, VisualBlock } from '@learnitive/project-semantic-index';
import type { ProjectFilePath } from '@learnitive/project-storage';

export const PROJECT_MEANING_SCHEMA_VERSION = 1 as const;
export const EMBEDDING_MODEL_VERSION = 'embeddinggemma-300m-onnx-q8-75a84c7-v1' as const;
export const EMBEDDING_MODEL_ID = 'embeddinggemma-300m-q8' as const;
export const EMBEDDING_DIMENSIONS = 768 as const;

export type EmbeddingModelId = typeof EMBEDDING_MODEL_ID;
export type EmbeddingModelStatus = 'not-installed' | 'downloading' | 'installed' | 'loading' | 'ready' | 'error';

export type EmbeddingModelState = {
  modelId: EmbeddingModelId;
  status: EmbeddingModelStatus;
  progress?: number;
  loadedBytes?: number;
  totalBytes?: number;
  error?: string;
};

export type EmbeddingCapabilities = {
  supported: boolean;
  webGpu: boolean;
  wasm: boolean;
  reason?: string;
};

export interface EmbeddingEngine {
  load(signal?: AbortSignal): Promise<void>;
  unload(): Promise<void>;
  embedDocuments(texts: readonly string[], signal?: AbortSignal): Promise<Float32Array[]>;
  embedQuery(query: string, signal?: AbortSignal): Promise<Float32Array>;
}

export type ProjectMeaningChunkKind =
  | 'prose'
  | 'heading'
  | 'theorem'
  | 'equation-context'
  | 'figure-caption'
  | 'table-caption'
  | 'list'
  | 'source'
  | 'bibliography';

export type ProjectMeaningChunk = {
  id: string;
  file: ProjectFilePath;
  range: SourceRange;
  line: number;
  kind: ProjectMeaningChunkKind;
  text: string;
  title?: string;
};

export type StoredProjectMeaningChunk = ProjectMeaningChunk & { contentHash: string };

export type StoredProjectMeaningIndex = {
  schemaVersion: typeof PROJECT_MEANING_SCHEMA_VERSION;
  generationId: string;
  modelId: EmbeddingModelId;
  modelVersion: typeof EMBEDDING_MODEL_VERSION;
  dimensions: number;
  updatedAt: number;
  chunks: readonly StoredProjectMeaningChunk[];
  vectors: ArrayBuffer;
};

export interface ProjectMeaningStore {
  load(memoryId: string): Promise<StoredProjectMeaningIndex | undefined>;
  save(memoryId: string, index: StoredProjectMeaningIndex): Promise<void>;
  clear(memoryId: string): Promise<void>;
}

export interface EmbeddingPlatform {
  readonly modelId: EmbeddingModelId;
  capabilities(): Promise<EmbeddingCapabilities>;
  modelState(): Promise<EmbeddingModelState>;
  prepareModel(onProgress?: (state: EmbeddingModelState) => void, signal?: AbortSignal): Promise<EmbeddingModelState>;
  removeModel(): Promise<void>;
  createEngine(): EmbeddingEngine;
  readonly store: ProjectMeaningStore;
  dispose(): void;
}

export type ProjectMeaningSearchResult = {
  chunk: StoredProjectMeaningChunk;
  score: number;
};

function compact(value: string): string {
  return value.replace(/\s+/gu, ' ').trim();
}

function lineForBlock(block: VisualBlock, snapshot: ProjectSemanticSnapshot): number {
  const file = snapshot.files.get(block.file);
  const node = file?.nodes.find((item) => item.range.from >= block.range.from && item.range.from < block.range.to);
  return node?.line ?? 1;
}

function blockText(block: VisualBlock): { kind: ProjectMeaningChunkKind; text: string; title?: string } | undefined {
  if (block.kind === 'heading') {
    const text = compact(block.text);
    return text ? { kind: 'heading', text, title: text } : undefined;
  }
  if (block.kind === 'paragraph' || block.kind === 'rich-paragraph' || block.kind === 'text-environment') {
    const text = compact(block.text);
    if (!text) return undefined;
    const equation = (block.kind === 'rich-paragraph' || block.kind === 'text-environment') && block.inlines.some((inline) => inline.kind === 'math');
    return { kind: equation ? 'equation-context' : 'prose', text };
  }
  if (block.kind === 'theorem') {
    const text = compact(block.text);
    return text ? { kind: 'theorem', text, ...(block.title ? { title: compact(block.title) } : {}) } : undefined;
  }
  if (block.kind === 'figure') {
    const text = compact(block.caption ?? block.text);
    return text ? { kind: 'figure-caption', text, ...(block.caption ? { title: compact(block.caption) } : {}) } : undefined;
  }
  if (block.kind === 'table') {
    const cells = block.cells?.map((cell) => compact(cell.text)).filter(Boolean).join(' | ') ?? '';
    const text = compact([block.caption, cells].filter(Boolean).join(' — ') || block.text);
    return text ? { kind: 'table-caption', text, ...(block.caption ? { title: compact(block.caption) } : {}) } : undefined;
  }
  if (block.kind === 'list') {
    const text = compact(block.items.map((item) => item.inlines.map((inline) => 'text' in inline ? inline.text : '').join(' ')).join(' '));
    return text ? { kind: 'list', text } : undefined;
  }
  if (block.kind === 'raw') {
    const text = compact(block.text);
    if (text.length < 40 || text.length > 4000) return undefined;
    return { kind: 'source', text };
  }
  return undefined;
}

/**
 * Builds stable, source-ranged project passages from the existing deterministic
 * semantic projection. It deliberately does not parse LaTeX a second time.
 */
export function buildProjectMeaningChunks(snapshot: ProjectSemanticSnapshot): ProjectMeaningChunk[] {
  const result: ProjectMeaningChunk[] = [];
  for (const [file, index] of [...snapshot.files.entries()].sort(([a], [b]) => a.localeCompare(b))) {
    let heading: string | undefined;
    let proseGroup: { from: number; to: number; line: number; texts: string[]; title?: string } | undefined;
    const flushProse = () => {
      if (!proseGroup) return;
      const text = compact(proseGroup.texts.join(' '));
      if (text) {
        result.push({
          id: `${file}:${proseGroup.from}:${proseGroup.to}:prose`,
          file,
          range: { from: proseGroup.from, to: proseGroup.to },
          line: proseGroup.line,
          kind: 'prose',
          text,
          ...(proseGroup.title ? { title: proseGroup.title } : {}),
        });
      }
      proseGroup = undefined;
    };

    for (const block of index.visualBlocks) {
      const mapped = blockText(block);
      if (!mapped) continue;
      if (mapped.kind === 'heading') {
        flushProse();
        heading = mapped.title;
        result.push({
          id: `${file}:${block.range.from}:${block.range.to}:heading`,
          file,
          range: block.range,
          line: lineForBlock(block, snapshot),
          kind: 'heading',
          text: mapped.text,
          ...(heading ? { title: heading } : {}),
        });
        continue;
      }
      if (mapped.kind === 'prose') {
        const line = lineForBlock(block, snapshot);
        if (!proseGroup) {
          proseGroup = { from: block.range.from, to: block.range.to, line, texts: [mapped.text], ...(heading ? { title: heading } : {}) };
        } else if (proseGroup.texts.join(' ').length + mapped.text.length <= 2600) {
          proseGroup.to = block.range.to;
          proseGroup.texts.push(mapped.text);
        } else {
          flushProse();
          proseGroup = { from: block.range.from, to: block.range.to, line, texts: [mapped.text], ...(heading ? { title: heading } : {}) };
        }
        continue;
      }
      flushProse();
      const contextualTitle = mapped.title ?? heading;
      result.push({
        id: `${file}:${block.range.from}:${block.range.to}:${mapped.kind}`,
        file,
        range: block.range,
        line: lineForBlock(block, snapshot),
        kind: mapped.kind,
        text: mapped.text,
        ...(contextualTitle ? { title: contextualTitle } : {}),
      });
    }
    flushProse();
  }

  for (const entry of snapshot.bibliographyEntries) {
    const field = (name: string) => entry.fields.find((item) => item.name.toLowerCase() === name)?.value.trim();
    const title = field('title');
    const text = compact([title, field('author'), field('year'), field('journal') ?? field('booktitle'), field('abstract')].filter(Boolean).join(' — '));
    if (!text) continue;
    result.push({
      id: `${entry.file}:${entry.range.from}:${entry.range.to}:bibliography`,
      file: entry.file,
      range: entry.range,
      line: entry.line,
      kind: 'bibliography',
      text,
      ...(title ? { title } : {}),
    });
  }
  return result;
}


export type ProjectMeaningReuseLookup = {
  byId: ReadonlyMap<string, { chunk: StoredProjectMeaningChunk; index: number }>;
  byHash: ReadonlyMap<string, number>;
};

export function createProjectMeaningReuseLookup(
  index: StoredProjectMeaningIndex,
): ProjectMeaningReuseLookup {
  const byId = new Map<string, { chunk: StoredProjectMeaningChunk; index: number }>();
  const byHash = new Map<string, number>();
  index.chunks.forEach((chunk, position) => {
    byId.set(chunk.id, { chunk, index: position });
    if (!byHash.has(chunk.contentHash)) byHash.set(chunk.contentHash, position);
  });
  return { byId, byHash };
}

export function reusableProjectMeaningVectorIndex(
  lookup: ProjectMeaningReuseLookup,
  chunk: StoredProjectMeaningChunk,
): number | undefined {
  const exact = lookup.byId.get(chunk.id);
  if (exact?.chunk.contentHash === chunk.contentHash) return exact.index;
  return lookup.byHash.get(chunk.contentHash);
}

export function normalizeVector(vector: Float32Array): Float32Array {
  let sum = 0;
  for (const value of vector) sum += value * value;
  const norm = Math.sqrt(sum);
  if (!Number.isFinite(norm) || norm <= 0) return vector.slice();
  const normalized = new Float32Array(vector.length);
  for (let i = 0; i < vector.length; i += 1) normalized[i] = vector[i]! / norm;
  return normalized;
}

export function cosineSimilarity(left: Float32Array, right: Float32Array): number {
  if (left.length !== right.length || left.length === 0) return Number.NEGATIVE_INFINITY;
  let dot = 0;
  let leftNorm = 0;
  let rightNorm = 0;
  for (let i = 0; i < left.length; i += 1) {
    const a = left[i]!;
    const b = right[i]!;
    dot += a * b;
    leftNorm += a * a;
    rightNorm += b * b;
  }
  const denominator = Math.sqrt(leftNorm) * Math.sqrt(rightNorm);
  return denominator > 0 ? dot / denominator : Number.NEGATIVE_INFINITY;
}

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

export function rankProjectMeaning(
  index: StoredProjectMeaningIndex,
  queryVector: Float32Array,
  limit = 30,
): ProjectMeaningSearchResult[] {
  if (queryVector.length !== index.dimensions) throw new Error('Embedding dimensions do not match the project index');
  return index.chunks
    .map((chunk, position) => ({ chunk, score: cosineSimilarity(queryVector, vectorAt(index, position)) }))
    .filter((item) => Number.isFinite(item.score))
    .sort((a, b) => b.score - a.score)
    .slice(0, Math.max(1, limit));
}