import { askLLM, pickModel } from './llm.mjs';
import { parseLLMJson } from './json-utils.mjs';
import { addKGEntity, addKGRelation, findKGEntities } from './knowledge-graph.mjs';
import { getEntityByName } from './db.mjs';

export async function extractEntities(text, opts = {}) {
  const maxLen = opts.maxLength || 6000;
  const prompt = `Extract key entities and relationships from this text.

Text:
${text.slice(0, maxLen)}

Return JSON only:
{
  "entities": [{ "name": "normalized name", "type": "person|technology|concept|organization|place|event|document|other", "metadata": "brief description" }],
  "relations": [{ "source": "entity name", "target": "entity name", "relation": "verb phrase like uses/created_by/related_to/part_of" }]
}
Rules:
- Max 20 entities
- Normalize names (lowercase, trimmed, no duplicates)
- Only extract clearly stated facts
- Entity types: person, technology, concept, organization, place, event, document, other`;

  try {
    const response = await askLLM(prompt, { model: pickModel('extract') });
    const parsed = parseLLMJson(response);
    if (!parsed.ok) return { entities: [], relations: [] };
    return {
      entities: (parsed.data.entities || []).map(e => ({
        name: String(e.name || '').trim().toLowerCase(),
        type: e.type || 'other',
        metadata: e.metadata || null
      })).filter(e => e.name),
      relations: (parsed.data.relations || []).map(r => ({
        source: String(r.source || '').trim().toLowerCase(),
        target: String(r.target || '').trim().toLowerCase(),
        relation: r.relation || 'related_to'
      })).filter(r => r.source && r.target)
    };
  } catch {
    return { entities: [], relations: [] };
  }
}

export async function extractAndStore(text, workspaceId = 'default', opts = {}) {
  const extraction = await extractEntities(text, opts);
  const entityIds = [];
  const entityNameToId = new Map();

  // Store entities
  for (const e of extraction.entities) {
    const id = addKGEntity(e.name, e.type, workspaceId, e.metadata);
    entityIds.push(id);
    entityNameToId.set(e.name, id);
  }

  // Store relations
  const relationIds = [];
  for (const r of extraction.relations) {
    let srcId = entityNameToId.get(r.source);
    let tgtId = entityNameToId.get(r.target);
    // Look up by name if not in current batch
    if (!srcId) { const e = getEntityByName(r.source, 'other', workspaceId); srcId = e?.id; }
    if (!tgtId) { const e = getEntityByName(r.target, 'other', workspaceId); tgtId = e?.id; }
    if (srcId && tgtId) {
      const relId = addKGRelation(srcId, tgtId, r.relation);
      relationIds.push(relId);
    }
  }

  return { entityIds, relationIds, raw: extraction };
}
