import { useCallback, useRef, useState } from 'react'
import { ReactFlow, Background, Controls, MiniMap, addEdge, useEdgesState, useNodesState, type Connection } from '@xyflow/react'
import { Download, Play, Plus, RotateCcw, Upload } from 'lucide-react'
import WorkflowNode from './WorkflowNode'
import type { ConnectionType, NodeData, NodeKind, RunResult, Workflow, WorkflowEdge, WorkflowNode as WNode } from './types'

const blank = (mode: NodeKind, label: string, extra: Partial<NodeData> = {}): NodeData => ({ label, mode, prompt: '', value: '', description: '', input_guidance: '', ...extra })
const initialNodes: WNode[] = [
  { id: 'input-1', type: 'workflow', position: { x: 20, y: 180 }, data: blank('Input', 'Customer message', { value: 'My order arrived damaged. Can you help?' }) },
  { id: 'worker-1', type: 'workflow', position: { x: 300, y: 180 }, data: blank('AI Worker', 'Support agent', { prompt: 'Investigate the request when useful, then write a concise recommended response.' }) },
  { id: 'tool-1', type: 'workflow', position: { x: 310, y: 390 }, data: blank('Simulated Tool', 'Customer record lookup', { description: 'Use this to retrieve a customer profile, order history, and support history.', input_guidance: 'Customer details or clues from the current message.', prompt: 'Generate a plausible CRM record. Include account tier, recent order, and relevant support history.' }) },
  { id: 'decision-1', type: 'workflow', position: { x: 590, y: 180 }, data: blank('Decision Rules', 'Risk check', { prompt: 'If the proposed response mentions a refund above $100, legal risk, or safety risk, choose Human review. Otherwise choose Final result.' }) },
  { id: 'review-1', type: 'workflow', position: { x: 850, y: 80 }, data: blank('Human Review', 'Manager approval', { prompt: 'Review the proposed response. Approve it or explain what should change.' }) },
  { id: 'output-1', type: 'workflow', position: { x: 850, y: 280 }, data: blank('Output', 'Final result') },
]
const flow = (id: string, source: string, target: string, label = ''): WorkflowEdge => ({ id, source, target, sourceHandle: 'flow-out', targetHandle: 'flow-in', label, data: { connectionType: 'flow' } })
const initialEdges: WorkflowEdge[] = [
  flow('e1', 'input-1', 'worker-1'), flow('e2', 'worker-1', 'decision-1'), flow('e3', 'decision-1', 'review-1', 'Human review'), flow('e4', 'decision-1', 'output-1', 'Final result'), flow('e5', 'review-1', 'output-1'),
  { id: 'e-tool', source: 'worker-1', target: 'tool-1', sourceHandle: 'tool-out', targetHandle: 'tool-in', label: 'can call', data: { connectionType: 'tool' }, animated: true, style: { stroke: '#9a62d5', strokeDasharray: '6 5' } },
]
const nodeTypes = { workflow: WorkflowNode }
const copy: Record<NodeKind, { label: string; help: string }> = {
  Input: { label: 'Input', help: 'One item enters' }, 'AI Worker': { label: 'AI Worker', help: 'Reasons and calls tools' }, 'Decision Rules': { label: 'Decision Rules', help: 'Routes using English rules' },
  'Human Review': { label: 'Human Review', help: 'Pauses for a person' }, 'Simulated Tool': { label: 'Simulated Tool', help: 'Returns simulated information or actions' }, Output: { label: 'Output', help: 'The run finishes' },
}

export default function App() {
  const [nodes, setNodes, onNodesChange] = useNodesState<WNode>(initialNodes)
  const [edges, setEdges, onEdgesChange] = useEdgesState<WorkflowEdge>(initialEdges)
  const [selected, setSelected] = useState<string | null>('worker-1')
  const [selectedEdge, setSelectedEdge] = useState<string | null>(null)
  const [name, setName] = useState('Agentic customer support')
  const [result, setResult] = useState<RunResult | null>(null)
  const [review, setReview] = useState('Approve and continue.')
  const [busy, setBusy] = useState(false)
  const [error, setError] = useState('')
  const fileRef = useRef<HTMLInputElement>(null)
  const node = nodes.find(n => n.id === selected)
  const edge = edges.find(e => e.id === selectedEdge)

  const onConnect = useCallback((connection: Connection) => {
    const type: ConnectionType = connection.sourceHandle === 'tool-out' || connection.targetHandle === 'tool-in' ? 'tool' : 'flow'
    const source = nodes.find(n => n.id === connection.source), target = nodes.find(n => n.id === connection.target)
    if (type === 'tool' && (source?.data.mode !== 'AI Worker' || target?.data.mode !== 'Simulated Tool')) { setError('Tool connections must go from an AI Worker’s purple tool handle to a Simulated Tool.'); return }
    if (type === 'flow' && (source?.data.mode === 'Simulated Tool' || target?.data.mode === 'Simulated Tool')) { setError('Simulated Tools use purple dashed tool connections, not workflow connections.'); return }
    setError('')
    setEdges(es => addEdge({ ...connection, label: type === 'tool' ? 'can call' : '', data: { connectionType: type }, animated: type === 'tool', style: type === 'tool' ? { stroke: '#9a62d5', strokeDasharray: '6 5' } : undefined }, es))
  }, [nodes, setEdges])

  const patchNode = (patch: Partial<NodeData>) => setNodes(ns => ns.map(n => n.id === selected ? { ...n, data: { ...n.data, ...patch } } : n))
  const addNode = (mode: NodeKind) => {
    const id = `${mode.toLowerCase().replaceAll(' ', '-')}-${Date.now()}`
    const defaults: Partial<NodeData> = mode === 'AI Worker' ? { prompt: 'Complete the task. Call an available tool only when its description is relevant.' } : mode === 'Decision Rules' ? { prompt: 'Write explicit routing rules here.' } : mode === 'Human Review' ? { prompt: 'Review the current item and say what should happen next.' } : mode === 'Simulated Tool' ? { description: 'Explain when an AI Worker should call this tool.', input_guidance: 'Describe the input the tool expects.', prompt: 'Describe how to generate the simulated output.' } : {}
    setNodes(ns => [...ns, { id, type: 'workflow', position: { x: 240 + Math.random() * 420, y: 80 + Math.random() * 320 }, data: blank(mode, `New ${copy[mode].label.toLowerCase()}`, defaults) }])
    setSelected(id); setSelectedEdge(null)
  }
  const workflow = (): Workflow => ({ name, nodes, edges })
  const call = async (url: string, body: object) => { setBusy(true); setError(''); try { const response = await fetch(url, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(body) }); const data = await response.json(); if (!response.ok) throw new Error(data.detail || 'The run failed.'); setResult(data) } catch (e) { setError(e instanceof Error ? e.message : 'The run failed.') } finally { setBusy(false) } }
  const download = () => { const blob = new Blob([JSON.stringify(workflow(), null, 2)], { type: 'application/json' }); const a = document.createElement('a'); a.href = URL.createObjectURL(blob); a.download = `${name.toLowerCase().replace(/[^a-z0-9]+/g, '-') || 'workflow'}.json`; a.click(); URL.revokeObjectURL(a.href) }
  const load = async (file?: File) => { if (!file) return; try { const data = JSON.parse(await file.text()) as Workflow; setName(data.name); setNodes(data.nodes); setEdges(data.edges.map(e => ({ ...e, data: e.data || { connectionType: 'flow' } }))); setSelected(null); setSelectedEdge(null); setResult(null); setError('') } catch { setError('That file is not valid workflow JSON.') } }

  return <main>
    <header><div><span className="eyebrow">EDUCATIONAL SANDBOX</span><h1>Agentic Workflow Lab</h1></div><div className="header-actions"><button className="quiet" onClick={() => fileRef.current?.click()}><Upload size={16}/> Load</button><input hidden ref={fileRef} type="file" accept="application/json" onChange={e => load(e.target.files?.[0])}/><button className="quiet" onClick={download}><Download size={16}/> Save</button><button className="run" disabled={busy} onClick={() => call('/api/run', { workflow: workflow(), input: nodes.find(n => n.data.mode === 'Input')?.data.value || '' })}><Play size={16}/>{busy ? 'Running…' : 'Run one item'}</button></div></header>
    <div className="workspace">
      <aside className="left-panel"><label>Workflow name<input value={name} onChange={e => setName(e.target.value)}/></label><h2>Building blocks</h2>{(Object.keys(copy) as NodeKind[]).map(mode => <button key={mode} className={`add-card ${mode.toLowerCase().replaceAll(' ', '-')}-card`} onClick={() => addNode(mode)}><Plus size={17}/><span><b>{copy[mode].label}</b><small>{copy[mode].help}</small></span></button>)}<div className="legend"><b>Connections</b><span><i className="flow-line"/> Workflow flow</span><span><i className="tool-line"/> Tool access</span></div><p className="hint">Connect solid side handles for workflow flow. Connect the AI Worker’s purple bottom handle to a tool’s purple top handle to grant tool access.</p></aside>
      <section className="canvas"><ReactFlow nodes={nodes} edges={edges} onNodesChange={onNodesChange} onEdgesChange={onEdgesChange} onConnect={onConnect} nodeTypes={nodeTypes} onNodeClick={(_, n) => { setSelected(n.id); setSelectedEdge(null) }} onEdgeClick={(_, e) => { setSelectedEdge(e.id); setSelected(null) }} onPaneClick={() => { setSelected(null); setSelectedEdge(null) }} fitView deleteKeyCode={['Backspace','Delete']}><Background gap={24} color="#d9ded8"/><Controls/><MiniMap pannable zoomable nodeColor={n => ({ Input:'#e7af2c','AI Worker':'#7975e8','Decision Rules':'#ea8d50','Human Review':'#e36e78','Simulated Tool':'#9a62d5',Output:'#2eb881' }[n.data.mode as NodeKind])}/></ReactFlow></section>
      <aside className="right-panel">{edge ? <><span className="eyebrow">SELECTED CONNECTION</span><h2>{edge.data?.connectionType === 'tool' ? 'Tool access' : 'Workflow flow'}</h2><p className="panel-copy">{edge.data?.connectionType === 'tool' ? 'This AI Worker may call the connected tool when its description is relevant.' : 'The item advances along this connection.'}</p>{edge.data?.connectionType === 'flow' && <label>Route label<input value={String(edge.label || '')} placeholder="e.g. High risk" onChange={e => setEdges(es => es.map(x => x.id === edge.id ? { ...x, label: e.target.value } : x))}/></label>}<button className="danger" onClick={() => { setEdges(es => es.filter(e => e.id !== edge.id)); setSelectedEdge(null) }}>Delete connection</button></> : node ? <><span className="eyebrow">SELECTED BLOCK</span><h2>{node.data.mode}</h2><label>Name<input value={node.data.label} onChange={e => patchNode({ label: e.target.value })}/></label>{node.data.mode === 'Input' && <label>Item<textarea rows={7} value={node.data.value} onChange={e => patchNode({ value: e.target.value })}/></label>}{['AI Worker','Decision Rules','Human Review'].includes(node.data.mode) && <label>{node.data.mode === 'Decision Rules' ? 'Routing rules' : node.data.mode === 'Human Review' ? 'Review request' : 'Worker instructions'}<textarea rows={8} value={node.data.prompt} onChange={e => patchNode({ prompt: e.target.value })}/></label>}{node.data.mode === 'Simulated Tool' && <><label>Tool description<textarea rows={4} value={node.data.description} onChange={e => patchNode({ description: e.target.value })}/><small className="field-help">The AI Worker uses this to decide when to call the tool.</small></label><label>Expected input<textarea rows={3} value={node.data.input_guidance} onChange={e => patchNode({ input_guidance: e.target.value })}/></label><label>Simulation instructions<textarea rows={6} value={node.data.prompt} onChange={e => patchNode({ prompt: e.target.value })}/></label></>}<button className="danger" onClick={() => { setNodes(ns => ns.filter(n => n.id !== node.id)); setEdges(es => es.filter(e => e.source !== node.id && e.target !== node.id)); setSelected(null) }}>Delete block</button></> : <div className="empty"><RotateCcw size={24}/><h2>Select something</h2><p>Edit a block or label a Decision Rules route.</p></div>}</aside>
    </div>
    <section className={`run-drawer ${result || error ? 'open' : ''}`}><div><span className="eyebrow">LATEST RUN</span><h2>{error ? 'Needs attention' : result?.status === 'awaiting_review' ? 'Paused for human review' : 'Run complete'}</h2></div>{error && <p className="error">{error}</p>}{result?.status === 'awaiting_review' && <div className="review"><p>{result.message}</p><blockquote>{result.result}</blockquote><textarea rows={3} value={review} onChange={e => setReview(e.target.value)}/><button className="run" disabled={busy} onClick={() => call('/api/resume', { run_id: result.run_id, response: review })}>Continue workflow</button></div>}{result?.status === 'completed' && <><pre>{result.result}</pre><details><summary>View {result.trace.length} execution events</summary>{result.trace.map((t,i) => <div className={`trace ${t.kind === 'tool' ? 'tool-event' : ''}`} key={i}><b>{i+1}. {t.label}</b><span>{t.kind === 'tool' ? 'TOOL CALL' : t.mode}</span><p>{t.output}</p></div>)}</details></>}</section>
  </main>
}