import { useCallback, useEffect, useRef, useState } from 'react';
import * as d3 from 'd3';
import { Trash2, Edit2, Save, X, RefreshCw, AlertCircle, Plus, Link as LinkIcon } from 'lucide-react';
import * as db from '../../lib/db';
import { callLLM } from '../../lib/llm';
import { EMIT_GRAPH_ACTIONS_TOOL } from '../../lib/llmTools';
import Button from '../ui/Button';
import SuggestionsQueue from './SuggestionsQueue';

const KnowledgeGraph = () => {
  const svgRef = useRef(null);
  const wrapperRef = useRef(null);
  const [dimensions, setDimensions] = useState({ width: 800, height: 600 });
  const [selectedNode, setSelectedNode] = useState(null);
  const [isEditing, setIsEditing] = useState(false);
  const [editData, setEditData] = useState({});
  const [isAnalyzing, setIsAnalyzing] = useState(false);
  const [analyzeError, setAnalyzeError] = useState(null);
  const [newRelData, setNewRelData] = useState({ target: '', type: 'related_to' });


  const [topics, setTopics] = useState([]);
  const [relations, setRelations] = useState([]);
  const [showExcludeNodes, setShowExcludeNodes] = useState(false);

  const reloadKb = useCallback(() => {
    Promise.all([db.getTopics(), db.getRelations()]).then(([allTopics, allRelations]) => {
      const topicIds = new Set(allTopics.map(t => t.id));
      const validRelations = allRelations
        .map(r => ({
          ...r,
          source: r.source?.id || r.source,
          target: r.target?.id || r.target
        }))
        .filter(r => topicIds.has(r.source) && topicIds.has(r.target));
      
      setTopics(allTopics);
      setRelations(validRelations);
    });
  }, []);

  useEffect(() => {
    reloadKb();
    const handler = () => reloadKb();
    window.addEventListener('respellion:kb-updated', handler);
    return () => window.removeEventListener('respellion:kb-updated', handler);
  }, [reloadKb]);

  useEffect(() => {
    if (!wrapperRef.current) return;
    const { width, height } = wrapperRef.current.getBoundingClientRect();
    setDimensions({ width, height });

    const handleResize = () => {
      if (wrapperRef.current) {
        const { width, height } = wrapperRef.current.getBoundingClientRect();
        setDimensions({ width, height });
      }
    };
    window.addEventListener('resize', handleResize);
    return () => window.removeEventListener('resize', handleResize);
  }, []);

  useEffect(() => {
    if (!svgRef.current || topics.length === 0) return;

    const { width, height } = dimensions;
    d3.select(svgRef.current).selectAll('*').remove();

    const filteredTopics = showExcludeNodes ? topics : topics.filter(t => t.learning_relevance !== 'exclude');
    const filteredTopicIds = new Set(filteredTopics.map(t => t.id));
    const filteredRelations = relations.filter(r => 
      filteredTopicIds.has(r.source?.id || r.source) && 
      filteredTopicIds.has(r.target?.id || r.target)
    );

    const nodes = filteredTopics.map(d => ({ ...d }));
    const links = filteredRelations.map(d => ({ ...d }));

    const svg = d3.select(svgRef.current)
      .attr('viewBox', [0, 0, width, height])
      .attr('width', width)
      .attr('height', height);

    const g = svg.append('g');

    const zoom = d3.zoom()
      .scaleExtent([0.1, 4])
      .on('zoom', (event) => { g.attr('transform', event.transform); });

    svg.call(zoom);

    const color = d3.scaleOrdinal()
      .domain(['concept', 'role', 'process'])
      .range(['#1F5560', '#5C1DD8', '#AFC0FF']);

    const simulation = d3.forceSimulation(nodes)
      .force('link', d3.forceLink(links).id(d => d.id).distance(100))
      .force('charge', d3.forceManyBody().strength(-300))
      .force('center', d3.forceCenter(width / 2, height / 2))
      .force('collide', d3.forceCollide().radius(40));

    const link = g.append('g')
      .attr('stroke', 'var(--color-bg-warm)')
      .attr('stroke-opacity', 0.6)
      .selectAll('line')
      .data(links)
      .join('line')
      .attr('stroke-width', 2);

    const node = g.append('g')
      .selectAll('g')
      .data(nodes)
      .join('g')
      .call(drag(simulation));

    node.append('circle')
      .attr('r', 20)
      .attr('fill', d => color(d.type) || '#1F5560')
      .attr('stroke', '#fff')
      .attr('stroke-width', 2)
      .attr('stroke-dasharray', d => d.learning_relevance === 'exclude' ? '4,4' : 'none')
      .attr('opacity', d => {
        if (d.learning_relevance === 'exclude') return 0.3;
        if (d.learning_relevance === 'peripheral') return 0.5;
        return 1;
      })
      .on('click', (event, d) => {
        setSelectedNode(d);
        setIsEditing(false);
      });

    node.append('text')
      .text(d => d.label)
      .attr('x', 25)
      .attr('y', 5)
      .attr('font-size', '12px')
      .attr('fill', 'var(--color-fg)')
      .attr('class', 'font-mono select-none pointer-events-none');

    simulation.on('tick', () => {
      link
        .attr('x1', d => d.source.x)
        .attr('y1', d => d.source.y)
        .attr('x2', d => d.target.x)
        .attr('y2', d => d.target.y);
      node.attr('transform', d => `translate(${d.x},${d.y})`);
    });

    function drag(simulation) {
      function dragstarted(event) {
        if (!event.active) simulation.alphaTarget(0.3).restart();
        event.subject.fx = event.subject.x;
        event.subject.fy = event.subject.y;
      }
      function dragged(event) {
        event.subject.fx = event.x;
        event.subject.fy = event.y;
      }
      function dragended(event) {
        if (!event.active) simulation.alphaTarget(0);
        event.subject.fx = null;
        event.subject.fy = null;
      }
      return d3.drag().on('start', dragstarted).on('drag', dragged).on('end', dragended);
    }

    return () => simulation.stop();
  }, [dimensions, topics, relations]);

  const startEdit = () => {
    setEditData({ label: selectedNode.label, type: selectedNode.type, description: selectedNode.description });
    setIsEditing(true);
  };

  const saveEdit = async () => {
    // If the admin touched learning_relevance, lock it so re-extraction won't overwrite the choice.
    // But an explicit relevance_locked in editData (the unlock checkbox) always wins.
    const relevanceChanged =
      editData.learning_relevance !== undefined &&
      editData.learning_relevance !== selectedNode.learning_relevance;
    const next = { ...editData };
    if (relevanceChanged && next.relevance_locked === undefined) next.relevance_locked = true;
    await db.upsertTopic({ ...selectedNode, ...next });
    const updated = topics.map(t => t.id === selectedNode.id ? { ...t, ...next } : t);
    setTopics(updated);
    setSelectedNode({ ...selectedNode, ...next });
    setIsEditing(false);
  };

  const deleteNode = async () => {
    if (confirm(`Are you sure you want to delete "${selectedNode.label}"? This will also remove any relations connected to it.`)) {
      const updatedTopics = topics.filter(t => t.id !== selectedNode.id);
      const updatedRelations = relations
        .map(r => ({
          ...r,
          source: r.source?.id || r.source,
          target: r.target?.id || r.target
        }))
        .filter(r => r.source !== selectedNode.id && r.target !== selectedNode.id);

      await db.saveTopics(updatedTopics);
      await db.saveRelations(updatedRelations);
      await db.deleteContent(selectedNode.id);
      await db.setQuizBank(selectedNode.id, []);

      setTopics(updatedTopics);
      setRelations(updatedRelations);
      setSelectedNode(null);
      setIsEditing(false);
    }
  };

  const addRelation = async () => {
    if (!newRelData.target) return;
    const isDup = relations.some(r => (r.source.id || r.source) === selectedNode.id && (r.target.id || r.target) === newRelData.target && r.type === newRelData.type);
    if (isDup) return;

    const newRel = { source: selectedNode.id, target: newRelData.target, type: newRelData.type };
    await db.addRelation(newRel);
    setRelations([...relations, newRel]);
    setNewRelData({ target: '', type: 'related_to' });
  };

  const removeRelation = async (sourceId, targetId, type) => {
    await db.removeRelation(sourceId, targetId, type);
    setRelations(relations.filter(r => !(
      (r.source.id || r.source) === sourceId &&
      (r.target.id || r.target) === targetId &&
      r.type === type
    )));
  };


  const analyzeGraph = async () => {
    if (!confirm('This will send the entire graph to the AI to evaluate content, merge duplicates, and create new logical relations. This may take a moment. Proceed?')) return;

    setIsAnalyzing(true);
    setAnalyzeError(null);

    try {
      const currentTopics = await db.getTopics();
      const currentRelations = await db.getRelations();

      const systemPrompt = `You are a strict Data Quality AI maintaining a Knowledge Graph for Respellion.
Evaluate the provided topics and relations and emit the actions to take via the emit_graph_actions tool.

Rules:
1. Identify topics that mean exactly the same thing. Choose one to keep, one to delete (merges).
2. Identify topics that are too vague, irrelevant, or malformed (deletions).
3. Identify missing logical relations (depends_on, part_of, related_to, executed_by) between conceptually linked topics (newRelations).
4. Evaluate learning_relevance. Mark purely operational topics (printer guides, etc.) as "exclude"; low-priority as "peripheral" (relevanceUpdates).

Do not return the entire graph — only the actions to take.`;

      // Send a compact representation to minimize token usage and avoid rate limits.
      const compactTopics = currentTopics.map(({ id, label, type, learning_relevance }) => ({ id, label, type, learning_relevance }));
      const compactRelations = currentRelations.map(r => ({
        source: r.source?.id || r.source,
        target: r.target?.id || r.target,
        type: r.type
      }));

      const userPrompt = `Here is the current knowledge graph:
${JSON.stringify({ topics: compactTopics, relations: compactRelations })}`;

      const llmResult = await callLLM({
        task: 'graph.analyze',
        tier: 'reasoning',
        system: [{ type: 'text', text: systemPrompt, cache_control: { type: 'ephemeral' } }],
        user: userPrompt,
        tools: [EMIT_GRAPH_ACTIONS_TOOL],
        toolChoice: { type: 'tool', name: EMIT_GRAPH_ACTIONS_TOOL.name },
        maxTokens: 4096,
      });

      const actions = llmResult.toolUses[0]?.input;
      if (!actions) throw new Error('Graph analysis did not emit a tool result.');

      let updatedTopics = [...currentTopics];
      let updatedRelations = [...currentRelations];

      if (actions.merges && Array.isArray(actions.merges)) {
        for (const merge of actions.merges) {
          updatedTopics = updatedTopics.filter(t => t.id !== merge.deleteId);
          updatedRelations = updatedRelations.map(r => {
            if (r.source === merge.deleteId) return { ...r, source: merge.keepId };
            if (r.target === merge.deleteId) return { ...r, target: merge.keepId };
            return r;
          });
        }
      }

      if (actions.deletions && Array.isArray(actions.deletions)) {
        updatedTopics = updatedTopics.filter(t => !actions.deletions.includes(t.id));
        updatedRelations = updatedRelations.filter(r => !actions.deletions.includes(r.source) && !actions.deletions.includes(r.target));
      }

      if (actions.newRelations && Array.isArray(actions.newRelations)) {
        for (const newRel of actions.newRelations) {
          const isDup = updatedRelations.some(r => r.source === newRel.source && r.target === newRel.target && r.type === newRel.type);
          if (!isDup) updatedRelations.push(newRel);
        }
      }

      if (actions.relevanceUpdates && Array.isArray(actions.relevanceUpdates)) {
        for (const update of actions.relevanceUpdates) {
          const topicIndex = updatedTopics.findIndex(t => t.id === update.id);
          if (topicIndex !== -1 && !updatedTopics[topicIndex].relevance_locked) {
            updatedTopics[topicIndex] = { ...updatedTopics[topicIndex], learning_relevance: update.learning_relevance };
          }
        }
      }

      // Ensure all relations only reference existing nodes and are normalized to string IDs
      const finalTopicIds = new Set(updatedTopics.map(t => t.id));
      updatedRelations = updatedRelations
        .map(r => ({
          ...r,
          source: r.source?.id || r.source,
          target: r.target?.id || r.target
        }))
        .filter(r => 
          r.source !== r.target && 
          finalTopicIds.has(r.source) && 
          finalTopicIds.has(r.target)
        );

      await db.saveTopics(updatedTopics);
      await db.saveRelations(updatedRelations);

      setTopics(updatedTopics);
      setRelations(updatedRelations);
      setSelectedNode(null);

    } catch (e) {
      setAnalyzeError(e.message || 'Analysis failed. Please try again.');
    } finally {
      setIsAnalyzing(false);
    }
  };

  return (
    <div className="relative w-full h-full flex flex-col md:flex-row">
      <div ref={wrapperRef} className="flex-1 h-[400px] md:h-full cursor-grab active:cursor-grabbing border-r border-bg-warm">
        {topics.length === 0 ? (
          <div className="h-full flex items-center justify-center text-fg-muted p-8 text-center">
            No knowledge graph data yet. Upload source material in the Sources tab first.
          </div>
        ) : (
          <svg ref={svgRef} className="w-full h-full" />
        )}
      </div>

      <div className="w-full md:w-80 bg-paper p-6 overflow-y-auto border-t md:border-t-0 border-bg-warm flex flex-col">
        <div className="mb-6 pb-4 border-b border-bg-warm space-y-3">
          <label className="flex items-center gap-2 text-sm text-fg-muted cursor-pointer mb-3">
            <input 
              type="checkbox" 
              checked={showExcludeNodes} 
              onChange={e => setShowExcludeNodes(e.target.checked)}
              className="rounded bg-bg-warm border-transparent focus:ring-0 text-teal"
            />
            Show Excluded Nodes (Reference Material)
          </label>
          <div>
            <Button
              onClick={analyzeGraph}
              disabled={isAnalyzing || topics.length === 0}
              className="w-full flex justify-center items-center gap-2"
            >
              <RefreshCw size={16} className={isAnalyzing ? "animate-spin" : ""} />
              {isAnalyzing ? 'Analyzing Graph...' : 'Analyze & Optimize Graph'}
            </Button>
            {analyzeError && (
              <p className="mt-2 text-xs text-red-600 flex items-start gap-1">
                <AlertCircle size={14} className="shrink-0 mt-0.5" />
                {analyzeError}
              </p>
            )}
          </div>
        </div>

        {/* R42 chatbot suggestions queue */}
        <div className="mb-6 pb-4 border-b border-bg-warm">
          <SuggestionsQueue onApplied={reloadKb} />
        </div>

        <div className="flex items-center justify-between mb-4">
          <h3 className="font-bold text-lg text-teal">Node Details</h3>
          {selectedNode && !isEditing && (
            <div className="flex items-center gap-1">
              <button onClick={startEdit} className="p-1.5 text-fg-muted hover:text-teal rounded-[var(--r-sm)]" title="Edit">
                <Edit2 size={16} />
              </button>
              <button onClick={deleteNode} className="p-1.5 text-fg-muted hover:text-red-500 rounded-[var(--r-sm)]" title="Delete">
                <Trash2 size={16} />
              </button>
            </div>
          )}
        </div>

        {selectedNode ? (
          isEditing ? (
            <div className="space-y-4 flex-1">
              <div>
                <label className="text-xs text-fg-muted uppercase tracking-wider mb-1 block">Label</label>
                <input
                  className="w-full border border-bg-warm rounded-[var(--r-sm)] px-3 py-1.5 text-sm bg-bg"
                  value={editData.label}
                  onChange={e => setEditData({...editData, label: e.target.value})}
                />
              </div>
              <div>
                <label className="text-xs text-fg-muted uppercase tracking-wider mb-1 block">Type</label>
                <select
                  className="w-full border border-bg-warm rounded-[var(--r-sm)] px-3 py-1.5 text-sm bg-bg"
                  value={editData.type}
                  onChange={e => setEditData({...editData, type: e.target.value})}
                >
                  <option value="concept">Concept</option>
                  <option value="role">Role</option>
                  <option value="process">Process</option>
                </select>
              </div>
              <div>
                <label className="text-xs text-fg-muted uppercase tracking-wider mb-1 block">Learning Relevance</label>
                <select
                  className="w-full border border-bg-warm rounded-[var(--r-sm)] px-3 py-1.5 text-sm bg-bg"
                  value={editData.learning_relevance || 'standard'}
                  onChange={e => setEditData({...editData, learning_relevance: e.target.value})}
                >
                  <option value="core">Core</option>
                  <option value="standard">Standard</option>
                  <option value="peripheral">Peripheral</option>
                  <option value="exclude">Exclude</option>
                </select>
                {selectedNode.relevance_locked && (
                  <label className="flex items-center gap-2 text-xs text-fg-muted mt-2 cursor-pointer">
                    <input
                      type="checkbox"
                      checked={editData.relevance_locked !== false}
                      onChange={e => setEditData({...editData, relevance_locked: e.target.checked})}
                      className="rounded bg-bg-warm border-transparent focus:ring-0 text-teal"
                    />
                    Locked — re-extraction will not change this
                  </label>
                )}
              </div>
              <div>
                <label className="text-xs text-fg-muted uppercase tracking-wider mb-1 block">Description</label>
                <textarea
                  className="w-full border border-bg-warm rounded-[var(--r-sm)] px-3 py-1.5 text-sm bg-bg min-h-[100px]"
                  value={editData.description}
                  onChange={e => setEditData({...editData, description: e.target.value})}
                />
              </div>
              <div className="flex gap-2 pt-2">
                <Button onClick={saveEdit} className="flex-1 flex items-center justify-center gap-2"><Save size={16}/> Save</Button>
                <Button variant="outline" onClick={() => setIsEditing(false)} className="flex-1 flex items-center justify-center gap-2"><X size={16}/> Cancel</Button>
              </div>
            </div>
          ) : (
            <div className="space-y-4 flex-1">
              <div>
                <p className="text-xs text-fg-muted uppercase tracking-wider mb-1">Label</p>
                <p className="font-medium">{selectedNode.label}</p>
              </div>
              <div>
                <p className="text-xs text-fg-muted uppercase tracking-wider mb-1">Type & Relevance</p>
                <div className="flex gap-2 flex-wrap">
                  <span className="inline-block px-2 py-1 bg-bg-warm rounded-[var(--r-pill)] text-xs font-mono">{selectedNode.type}</span>
                  <span className="inline-block px-2 py-1 bg-bg-warm rounded-[var(--r-pill)] text-xs font-mono opacity-80">{selectedNode.learning_relevance || 'standard'}</span>
                  {selectedNode.relevance_locked && (
                    <span className="inline-block px-2 py-1 bg-teal/10 text-teal rounded-[var(--r-pill)] text-xs font-mono" title="Re-extraction will not change relevance for this topic.">locked</span>
                  )}
                </div>
              </div>
              <div>
                <p className="text-xs text-fg-muted uppercase tracking-wider mb-1">Description</p>
                <p className="text-sm leading-relaxed">{selectedNode.description}</p>
              </div>
              <div>
                <p className="text-xs text-fg-muted uppercase tracking-wider mb-1">ID</p>
                <p className="text-xs font-mono text-fg-muted break-all">{selectedNode.id}</p>
              </div>

              <div className="pt-4 border-t border-bg-warm">
                <p className="text-xs text-fg-muted uppercase tracking-wider mb-3">Relations</p>

                <div className="space-y-2 mb-4">
                  {relations.filter(r => (r.source.id || r.source) === selectedNode.id).map((rel, idx) => {
                    const targetId = rel.target.id || rel.target;
                    const targetTopic = topics.find(t => t.id === targetId);
                    return (
                      <div key={idx} className="flex items-center justify-between bg-bg rounded-[var(--r-sm)] p-2 border border-bg-warm">
                        <div className="text-xs">
                          <span className="font-mono text-teal">{rel.type}</span> &rarr; {targetTopic ? targetTopic.label : targetId}
                        </div>
                        <button
                          onClick={() => removeRelation(selectedNode.id, targetId, rel.type)}
                          className="text-fg-muted hover:text-red-500 p-1"
                          title="Remove Relation"
                        >
                          <X size={14} />
                        </button>
                      </div>
                    );
                  })}
                </div>

                <div className="bg-bg rounded-[var(--r-sm)] p-3 border border-bg-warm space-y-2">
                  <p className="text-xs font-medium flex items-center gap-1"><LinkIcon size={12}/> Add Relation</p>
                  <select
                    className="w-full border border-bg-warm rounded-[var(--r-sm)] px-2 py-1 text-xs bg-paper"
                    value={newRelData.type}
                    onChange={e => setNewRelData({...newRelData, type: e.target.value})}
                  >
                    <option value="related_to">related_to</option>
                    <option value="depends_on">depends_on</option>
                    <option value="part_of">part_of</option>
                    <option value="executed_by">executed_by</option>
                  </select>
                  <select
                    className="w-full border border-bg-warm rounded-[var(--r-sm)] px-2 py-1 text-xs bg-paper"
                    value={newRelData.target}
                    onChange={e => setNewRelData({...newRelData, target: e.target.value})}
                  >
                    <option value="">-- Select Target Topic --</option>
                    {topics.filter(t => t.id !== selectedNode.id).map(t => (
                      <option key={t.id} value={t.id}>{t.label}</option>
                    ))}
                  </select>
                  <Button
                    onClick={addRelation}
                    disabled={!newRelData.target}
                    className="w-full py-1 text-xs mt-1"
                  >
                    <Plus size={14} className="mr-1 inline" /> Add
                  </Button>
                </div>
              </div>
            </div>
          )
        ) : (
          <p className="text-sm text-fg-muted">Click a node in the graph to view or edit its details.</p>
        )}
      </div>
    </div>
  );
};

export default KnowledgeGraph;