import { NotFoundError } from '../../lib/errors.js';
import { getCollection } from '../../lib/datastore.js';
import type {
  KnowledgeBaseDoc,
  KnowledgeChunkDoc,
  KnowledgeSourceDoc,
  ListKnowledgeBasesQuery,
} from './types.js';

function baseCollection() {
  return getCollection<KnowledgeBaseDoc>('knowledge_bases', '/productId');
}

function sourceCollection() {
  return getCollection<KnowledgeSourceDoc>('knowledge_sources', '/knowledgeBaseId');
}

function chunkCollection() {
  return getCollection<KnowledgeChunkDoc>('knowledge_chunks', '/knowledgeBaseId');
}

export async function createBase(doc: KnowledgeBaseDoc): Promise<KnowledgeBaseDoc> {
  return baseCollection().create(doc);
}

export async function listBases(
  productId: string,
  query: ListKnowledgeBasesQuery
): Promise<KnowledgeBaseDoc[]> {
  return baseCollection().findMany({
    filter: {
      productId,
      ...(query.status ? { status: query.status } : {}),
    },
    sort: { createdAt: -1 },
    limit: query.limit,
  });
}

export async function getBase(id: string, productId: string): Promise<KnowledgeBaseDoc> {
  const base = await baseCollection().findById(id, productId);
  if (!base) throw new NotFoundError(`Knowledge base '${id}' not found`);
  return base;
}

export async function updateBase(
  id: string,
  productId: string,
  updates: Partial<KnowledgeBaseDoc>
): Promise<KnowledgeBaseDoc> {
  const updated = await baseCollection().update(id, productId, {
    ...updates,
    updatedAt: new Date().toISOString(),
  });
  if (!updated) throw new NotFoundError(`Knowledge base '${id}' not found`);
  return updated;
}

export async function createSource(doc: KnowledgeSourceDoc): Promise<KnowledgeSourceDoc> {
  return sourceCollection().create(doc);
}

export async function getSource(id: string, knowledgeBaseId: string): Promise<KnowledgeSourceDoc> {
  const source = await sourceCollection().findById(id, knowledgeBaseId);
  if (!source) throw new NotFoundError(`Knowledge source '${id}' not found`);
  return source;
}

export async function updateSource(
  id: string,
  knowledgeBaseId: string,
  updates: Partial<KnowledgeSourceDoc>
): Promise<KnowledgeSourceDoc> {
  const updated = await sourceCollection().update(id, knowledgeBaseId, {
    ...updates,
    updatedAt: new Date().toISOString(),
  });
  if (!updated) throw new NotFoundError(`Knowledge source '${id}' not found`);
  return updated;
}

export async function listSources(knowledgeBaseId: string): Promise<KnowledgeSourceDoc[]> {
  return sourceCollection().findMany({
    filter: { knowledgeBaseId },
    sort: { createdAt: -1 },
    limit: 200,
  });
}

export async function upsertChunks(docs: KnowledgeChunkDoc[]): Promise<KnowledgeChunkDoc[]> {
  const writes = docs.map(doc => chunkCollection().upsert(doc));
  return Promise.all(writes);
}

export async function listChunks(
  knowledgeBaseId: string,
  filter: { sourceId?: string; limit?: number } = {}
): Promise<KnowledgeChunkDoc[]> {
  return chunkCollection().findMany({
    filter: {
      knowledgeBaseId,
      ...(filter.sourceId ? { sourceId: filter.sourceId } : {}),
    },
    sort: { ordinal: 1 },
    limit: filter.limit ?? 500,
  });
}

export async function searchChunks(
  knowledgeBaseId: string,
  query: string,
  options: { sourceId?: string; limit?: number } = {}
): Promise<KnowledgeChunkDoc[]> {
  const allChunks = await listChunks(knowledgeBaseId, {
    sourceId: options.sourceId,
    limit: 1000,
  });

  const queryLower = query.toLowerCase();
  const terms = queryLower.split(/\s+/).filter(Boolean);

  const scored = allChunks
    .map(chunk => {
      const text = chunk.contentText.toLowerCase();
      let score = 0;
      for (const term of terms) {
        if (text.includes(term)) score += 1;
        if (chunk.tags.some(tag => tag.toLowerCase() === term)) score += 2;
      }
      return { chunk, score };
    })
    .filter(s => s.score > 0)
    .sort((a, b) => b.score - a.score)
    .slice(0, options.limit ?? 10);

  return scored.map(s => s.chunk);
}

export interface KnowledgeBaseStats {
  knowledgeBaseId: string;
  sourceCount: number;
  chunkCount: number;
  totalTokens: number;
  indexedSources: number;
  pendingSources: number;
  failedSources: number;
}

export async function getBaseStats(knowledgeBaseId: string): Promise<KnowledgeBaseStats> {
  const sources = await listSources(knowledgeBaseId);
  const chunks = await listChunks(knowledgeBaseId, { limit: 10000 });

  return {
    knowledgeBaseId,
    sourceCount: sources.length,
    chunkCount: chunks.length,
    totalTokens: chunks.reduce((sum, c) => sum + (c.tokenCount ?? 0), 0),
    indexedSources: sources.filter(s => s.status === 'indexed').length,
    pendingSources: sources.filter(s => s.status === 'pending').length,
    failedSources: sources.filter(s => s.status === 'failed').length,
  };
}

export async function deleteSource(id: string, knowledgeBaseId: string): Promise<void> {
  await sourceCollection().delete(id, knowledgeBaseId);
}

export async function deleteBase(id: string, productId: string): Promise<void> {
  await baseCollection().delete(id, productId);
}
