// PRE Web GUI — Knowledge Graph
// Entity-relationship graph layered on top of the flat memory system.
// Extracts entities (people, projects, tools, systems, concepts) from
// memories and experiences, connects them with typed edges, and supports
// traversal queries to surface related context the flat search would miss.

const fs = require('fs');
const path = require('path');
const { PRE_DIR } = require('./constants');
const { streamChat, embed } = require('./ollama');

const GRAPH_FILE = path.join(PRE_DIR, 'knowledge-graph.json');

// ── Graph schema ─────────────────────────────────────────────────────────────

const ENTITY_TYPES = [
  'person', 'project', 'system', 'tool', 'concept',
  'organization', 'location', 'file', 'model',
];

const RELATION_TYPES = [
  'works_on',     // person → project
  'uses',         // person|project → tool|system
  'depends_on',   // project|system → project|system
  'part_of',      // entity → organization|project
  'related_to',   // generic bidirectional
  'created_by',   // entity → person
  'located_at',   // entity → location
  'runs_on',      // system|model → system
  'trained_on',   // model → concept|file
  'mentions',     // memory/experience → entity (source provenance)
];

// ── Persistence ──────────────────────────────────────────────────────────────

let _graph = null;

function emptyGraph() {
  return {
    entities: {},     // id → { id, name, type, aliases, properties, sources, updatedAt }
    edges: [],        // [{ source, target, relation, weight, properties, createdAt }]
    version: 1,
    updatedAt: new Date().toISOString(),
  };
}

function loadGraph() {
  if (_graph) return _graph;
  try {
    if (fs.existsSync(GRAPH_FILE)) {
      _graph = JSON.parse(fs.readFileSync(GRAPH_FILE, 'utf-8'));
      // Ensure required fields
      if (!_graph.entities) _graph.entities = {};
      if (!_graph.edges) _graph.edges = [];
      return _graph;
    }
  } catch {}
  _graph = emptyGraph();
  return _graph;
}

function saveGraph() {
  if (!_graph) return;
  _graph.updatedAt = new Date().toISOString();
  try {
    fs.writeFileSync(GRAPH_FILE, JSON.stringify(_graph, null, 2), { mode: 0o600 });
  } catch (err) {
    console.log(`[knowledge-graph] Save error: ${err.message}`);
  }
}

// ── Entity operations ────────────────────────────────────────────────────────

function entityId(name) {
  return name.toLowerCase().replace(/[^a-z0-9]+/g, '_').replace(/^_|_$/g, '');
}

function addEntity(name, type, properties = {}, source = null) {
  const graph = loadGraph();
  const id = entityId(name);
  if (!ENTITY_TYPES.includes(type)) type = 'concept';

  if (graph.entities[id]) {
    // Merge — update properties, add alias if name differs
    const existing = graph.entities[id];
    if (existing.name !== name && !existing.aliases.includes(name)) {
      existing.aliases.push(name);
    }
    Object.assign(existing.properties, properties);
    if (source && !existing.sources.includes(source)) {
      existing.sources.push(source);
    }
    existing.updatedAt = new Date().toISOString();
  } else {
    graph.entities[id] = {
      id,
      name,
      type,
      aliases: [],
      properties,
      sources: source ? [source] : [],
      updatedAt: new Date().toISOString(),
    };
  }

  saveGraph();
  return graph.entities[id];
}

function removeEntity(nameOrId) {
  const graph = loadGraph();
  const id = entityId(nameOrId);
  if (!graph.entities[id]) return { error: `Entity '${nameOrId}' not found` };

  delete graph.entities[id];
  // Remove edges referencing this entity
  graph.edges = graph.edges.filter(e => e.source !== id && e.target !== id);
  saveGraph();
  return { success: true, removed: id };
}

function getEntity(nameOrId) {
  const graph = loadGraph();
  const id = entityId(nameOrId);
  return graph.entities[id] || null;
}

function listEntities(type = null) {
  const graph = loadGraph();
  const all = Object.values(graph.entities);
  if (type) return all.filter(e => e.type === type);
  return all;
}

function findEntity(query) {
  const graph = loadGraph();
  const q = query.toLowerCase();
  return Object.values(graph.entities).filter(e =>
    e.name.toLowerCase().includes(q) ||
    e.id.includes(q) ||
    e.aliases.some(a => a.toLowerCase().includes(q))
  );
}

// ── Edge operations ──────────────────────────────────────────────────────────

function addEdge(sourceName, targetName, relation, properties = {}) {
  const graph = loadGraph();
  const sourceId = entityId(sourceName);
  const targetId = entityId(targetName);

  if (!RELATION_TYPES.includes(relation)) relation = 'related_to';
  if (!graph.entities[sourceId] || !graph.entities[targetId]) {
    return { error: 'Both source and target entities must exist' };
  }

  // Check for duplicate
  const exists = graph.edges.some(e =>
    e.source === sourceId && e.target === targetId && e.relation === relation
  );
  if (exists) return { exists: true };

  graph.edges.push({
    source: sourceId,
    target: targetId,
    relation,
    weight: 1.0,
    properties,
    createdAt: new Date().toISOString(),
  });

  saveGraph();
  return { success: true };
}

function removeEdge(sourceName, targetName, relation) {
  const graph = loadGraph();
  const sourceId = entityId(sourceName);
  const targetId = entityId(targetName);

  const before = graph.edges.length;
  graph.edges = graph.edges.filter(e => {
    if (e.source !== sourceId || e.target !== targetId) return true;
    if (relation && e.relation !== relation) return true;
    return false;
  });

  if (graph.edges.length < before) {
    saveGraph();
    return { success: true, removed: before - graph.edges.length };
  }
  return { error: 'Edge not found' };
}

// ── Traversal ────────────────────────────────────────────────────────────────

/**
 * Get all entities directly connected to a given entity.
 * Returns { outgoing: [...], incoming: [...] } with relationship info.
 */
function getNeighbors(nameOrId) {
  const graph = loadGraph();
  const id = entityId(nameOrId);
  if (!graph.entities[id]) return { outgoing: [], incoming: [] };

  const outgoing = graph.edges
    .filter(e => e.source === id)
    .map(e => ({
      entity: graph.entities[e.target],
      relation: e.relation,
      weight: e.weight,
      properties: e.properties,
    }))
    .filter(n => n.entity);

  const incoming = graph.edges
    .filter(e => e.target === id)
    .map(e => ({
      entity: graph.entities[e.source],
      relation: e.relation,
      weight: e.weight,
      properties: e.properties,
    }))
    .filter(n => n.entity);

  return { outgoing, incoming };
}

/**
 * Get a subgraph around an entity within N hops.
 * Returns { entities: [...], edges: [...] }.
 */
function getSubgraph(nameOrId, maxHops = 2) {
  const graph = loadGraph();
  const startId = entityId(nameOrId);
  if (!graph.entities[startId]) return { entities: [], edges: [] };

  const visited = new Set([startId]);
  const frontier = [startId];
  const subEdges = [];

  for (let hop = 0; hop < maxHops && frontier.length > 0; hop++) {
    const nextFrontier = [];
    for (const nodeId of frontier) {
      for (const edge of graph.edges) {
        if (edge.source === nodeId && !visited.has(edge.target)) {
          visited.add(edge.target);
          nextFrontier.push(edge.target);
          subEdges.push(edge);
        } else if (edge.target === nodeId && !visited.has(edge.source)) {
          visited.add(edge.source);
          nextFrontier.push(edge.source);
          subEdges.push(edge);
        } else if (
          (edge.source === nodeId || edge.target === nodeId) &&
          !subEdges.includes(edge)
        ) {
          subEdges.push(edge);
        }
      }
    }
    frontier.length = 0;
    frontier.push(...nextFrontier);
  }

  const subEntities = [...visited]
    .map(id => graph.entities[id])
    .filter(Boolean);

  return { entities: subEntities, edges: subEdges };
}

/**
 * Search the graph for entities matching a query, then return their
 * neighborhood context — useful for enriching memory search results.
 */
function queryGraph(query) {
  const matches = findEntity(query);
  if (matches.length === 0) return { matches: [], context: '' };

  const contextParts = [];
  for (const entity of matches.slice(0, 5)) {
    const neighbors = getNeighbors(entity.id);
    const relations = [];

    for (const n of neighbors.outgoing) {
      relations.push(`${entity.name} —[${n.relation}]→ ${n.entity.name}`);
    }
    for (const n of neighbors.incoming) {
      relations.push(`${n.entity.name} —[${n.relation}]→ ${entity.name}`);
    }

    if (relations.length > 0) {
      contextParts.push(`${entity.name} (${entity.type}): ${relations.join('; ')}`);
    }
  }

  return {
    matches,
    context: contextParts.length > 0
      ? `<knowledge_graph>\n${contextParts.join('\n')}\n</knowledge_graph>`
      : '',
  };
}

// ── LLM-powered entity extraction ────────────────────────────────────────────

const EXTRACTION_PROMPT = `You are the knowledge graph extraction subagent for PRE (Personal Reasoning Engine).
Analyze the provided text and extract entities and relationships.

ENTITY TYPES: ${ENTITY_TYPES.join(', ')}
RELATION TYPES: ${RELATION_TYPES.join(', ')}

OUTPUT FORMAT — JSON only, no explanation:
{
  "entities": [{"name": "...", "type": "...", "properties": {}}],
  "edges": [{"source": "...", "target": "...", "relation": "..."}]
}

Rules:
- Extract only meaningful, reusable entities — not generic words
- Use the most specific entity type that fits
- Prefer existing relation types; use "related_to" as fallback
- Names should be proper nouns or specific identifiers, not descriptions
- Keep it concise: typically 2-8 entities and 1-6 edges per text
- If nothing meaningful can be extracted, return {"entities":[],"edges":[]}`;

/**
 * Extract entities and relationships from text using the LLM.
 * @param {string} text - Text to analyze (memory body, experience, conversation)
 * @param {string} source - Source identifier for provenance tracking
 * @returns {{ entities: Array, edges: Array }} Extracted graph elements
 */
async function extractFromText(text, source = null) {
  if (!text || text.length < 20) return { entities: [], edges: [] };

  try {
    const result = await streamChat({
      messages: [
        { role: 'system', content: EXTRACTION_PROMPT },
        { role: 'user', content: `Extract entities and relationships:\n\n${text.slice(0, 2000)}` },
      ],
      maxTokens: 512,
    });

    const response = result.response || '';
    const jsonMatch = response.match(/\{[\s\S]*\}/);
    if (!jsonMatch) return { entities: [], edges: [] };

    const parsed = JSON.parse(jsonMatch[0]);
    const entities = (parsed.entities || []).filter(e =>
      e.name && e.type && ENTITY_TYPES.includes(e.type)
    );
    const edges = (parsed.edges || []).filter(e =>
      e.source && e.target && e.relation
    );

    // Add entities to graph
    for (const ent of entities) {
      addEntity(ent.name, ent.type, ent.properties || {}, source);
    }

    // Add edges (entities must exist first)
    for (const edge of edges) {
      // Ensure both endpoints exist as entities
      const srcId = entityId(edge.source);
      const tgtId = entityId(edge.target);
      const graph = loadGraph();
      if (graph.entities[srcId] && graph.entities[tgtId]) {
        addEdge(edge.source, edge.target, edge.relation);
      }
    }

    return { entities, edges };
  } catch (err) {
    console.log(`[knowledge-graph] Extraction error: ${err.message}`);
    return { entities: [], edges: [] };
  }
}

/**
 * Build the graph from all existing memories and experiences.
 * Designed to run once for initial population, then incrementally.
 */
async function buildFromMemories() {
  const { getAllMemories } = require('./memory');
  const { listExperiences } = require('./experience');

  const memories = getAllMemories();
  const experiences = listExperiences();
  let extracted = 0;

  for (const mem of memories) {
    const text = `${mem.name}: ${mem.description}\n${mem.body}`;
    const result = await extractFromText(text, `memory:${mem.filename}`);
    extracted += result.entities.length;
  }

  for (const exp of experiences) {
    const text = `${exp.title || exp.name}: ${exp.body || exp.content || ''}`;
    const result = await extractFromText(text, `experience:${exp.filename}`);
    extracted += result.entities.length;
  }

  return { extracted, entities: Object.keys(loadGraph().entities).length, edges: loadGraph().edges.length };
}

// ── Auto-extraction integration ──────────────────────────────────────────────

// Throttling: extract graph entities less frequently than memory extraction
const graphExtractionState = {
  turnsSinceLastExtract: 0,
  running: false,
};

const GRAPH_EXTRACT_EVERY_N_TURNS = 6; // Every 6 turns (half as frequent as memory)

/**
 * Auto-extract entities from recent conversation messages.
 * Call this alongside memory autoExtract for incremental graph building.
 */
async function autoExtractGraph(messages) {
  graphExtractionState.turnsSinceLastExtract++;
  if (graphExtractionState.turnsSinceLastExtract < GRAPH_EXTRACT_EVERY_N_TURNS) return;
  if (graphExtractionState.running) return;

  graphExtractionState.running = true;
  graphExtractionState.turnsSinceLastExtract = 0;

  try {
    const recent = messages.slice(-6)
      .filter(m => m.role === 'user' || m.role === 'assistant')
      .map(m => `${m.role}: ${(m.content || '').slice(0, 500)}`)
      .join('\n');

    if (recent.length < 50) return;
    await extractFromText(recent, 'conversation');
  } catch (err) {
    console.log(`[knowledge-graph] Auto-extract error: ${err.message}`);
  } finally {
    graphExtractionState.running = false;
  }
}

// ── Context builder for system prompt ────────────────────────────────────────

/**
 * Build a concise knowledge graph context block for relevant entities.
 * Uses embedding similarity to find entities related to the user's query.
 */
async function buildGraphContext(userQuery) {
  if (!userQuery) return '';

  const graph = loadGraph();
  const entityCount = Object.keys(graph.entities).length;
  if (entityCount === 0) return '';

  // Try direct text matching first (fast path)
  const result = queryGraph(userQuery);
  if (result.context) return result.context;

  // Try embedding-based search for more nuanced matches
  try {
    const queryVectors = await embed(userQuery);
    if (!queryVectors || queryVectors.length === 0) return '';

    const queryVec = queryVectors[0];
    const entities = Object.values(graph.entities);

    // Embed entity names+types and find closest matches
    const entityTexts = entities.map(e =>
      `${e.name} (${e.type})${e.aliases.length ? ': ' + e.aliases.join(', ') : ''}`
    );
    const entityVectors = await embed(entityTexts);
    if (!entityVectors || entityVectors.length === 0) return '';

    // Score and rank
    const scored = entities.map((e, i) => {
      const vec = entityVectors[i];
      if (!vec) return { entity: e, score: 0 };
      let dot = 0, magA = 0, magB = 0;
      for (let j = 0; j < queryVec.length; j++) {
        dot += queryVec[j] * vec[j];
        magA += queryVec[j] * queryVec[j];
        magB += vec[j] * vec[j];
      }
      return { entity: e, score: dot / (Math.sqrt(magA) * Math.sqrt(magB) || 1) };
    });

    scored.sort((a, b) => b.score - a.score);
    const top = scored.filter(s => s.score > 0.3).slice(0, 5);
    if (top.length === 0) return '';

    const parts = [];
    for (const { entity } of top) {
      const neighbors = getNeighbors(entity.id);
      const rels = [];
      for (const n of neighbors.outgoing.slice(0, 3)) {
        rels.push(`→[${n.relation}] ${n.entity.name}`);
      }
      for (const n of neighbors.incoming.slice(0, 3)) {
        rels.push(`←[${n.relation}] ${n.entity.name}`);
      }
      const relStr = rels.length > 0 ? ` | ${rels.join(', ')}` : '';
      parts.push(`${entity.name} (${entity.type})${relStr}`);
    }

    return `<knowledge_graph>\n${parts.join('\n')}\n</knowledge_graph>`;
  } catch {
    return '';
  }
}

// ── Stats ────────────────────────────────────────────────────────────────────

function getStats() {
  const graph = loadGraph();
  const entities = Object.values(graph.entities);
  const typeCounts = {};
  for (const e of entities) {
    typeCounts[e.type] = (typeCounts[e.type] || 0) + 1;
  }

  const relationCounts = {};
  for (const edge of graph.edges) {
    relationCounts[edge.relation] = (relationCounts[edge.relation] || 0) + 1;
  }

  return {
    entityCount: entities.length,
    edgeCount: graph.edges.length,
    typeCounts,
    relationCounts,
    updatedAt: graph.updatedAt,
  };
}

module.exports = {
  // Entity operations
  addEntity,
  removeEntity,
  getEntity,
  listEntities,
  findEntity,

  // Edge operations
  addEdge,
  removeEdge,

  // Traversal
  getNeighbors,
  getSubgraph,
  queryGraph,

  // LLM extraction
  extractFromText,
  buildFromMemories,
  autoExtractGraph,

  // Context
  buildGraphContext,

  // Stats & persistence
  getStats,
  loadGraph,
  saveGraph,

  // Constants (for testing)
  ENTITY_TYPES,
  RELATION_TYPES,
  GRAPH_FILE,
  entityId,
};
