frontend/src/pages/GridViewPage.jsx

import { useState, useEffect, useCallback } from 'react'
import { useNavigate, useSearchParams } from 'react-router-dom'
import { useAuth } from '../contexts/AuthContext'
import api from '../services/api'
import ReactFlow, { 
  Background, 
  Controls, 
  MiniMap,
  useNodesState,
  useEdgesState,
  MarkerType,
  Handle,
  Position
} from 'reactflow'
import 'reactflow/dist/style.css'
import { 
  Zap, 
  RefreshCw, 
  Play, 
  AlertCircle,
  CheckCircle2,
  ToggleLeft,
  ToggleRight,
  Info,
  Download,
  Loader2,
  X,
  Plus,
  Power
} from 'lucide-react'

// Custom Bus Node with Handles so edges connect
function BusNode({ data }) {
  const isSlack = data.is_slack
  return (
    <div className={`px-3 py-2 rounded-lg shadow-lg border-2 relative ${
      isSlack 
        ? 'bg-gradient-to-br from-primary-500 to-primary-600 border-primary-400 text-white' 
        : 'bg-white border-dark-200 text-dark-800'
    }`}>
      <Handle type="target" position={Position.Top} style={{ background: '#94a3b8', width: 6, height: 6 }} />
      <Handle type="target" position={Position.Left} id="left-t" style={{ background: '#94a3b8', width: 6, height: 6 }} />
      <Handle type="source" position={Position.Bottom} style={{ background: '#94a3b8', width: 6, height: 6 }} />
      <Handle type="source" position={Position.Right} id="right-s" style={{ background: '#94a3b8', width: 6, height: 6 }} />
      <div className="text-xs font-medium">{data.label}</div>
      {data.load_mw > 0 && (
        <div className="text-[10px] opacity-75">
          {data.load_mw.toFixed(2)} MW
        </div>
      )}
    </div>
  )
}

const nodeTypes = {
  busNode: BusNode,
}

export default function GridViewPage() {
  const { user, token } = useAuth()
  const navigate = useNavigate()
  const [searchParams, setSearchParams] = useSearchParams()
  
  const [nodes, setNodes, onNodesChange] = useNodesState([])
  const [edges, setEdges, onEdgesChange] = useEdgesState([])
  const [loading, setLoading] = useState(true)
  const [optimizing, setOptimizing] = useState(false)
  const [simulating, setSimulating] = useState(false)
  const [downloadingReport, setDownloadingReport] = useState(false)
  const [error, setError] = useState(null)
  const [success, setSuccess] = useState(null)
  const [system, setSystem] = useState(searchParams.get('system') || 'case33bw')
  const [gridData, setGridData] = useState(null)
  const [powerFlow, setPowerFlow] = useState(null)
  const [openLines, setOpenLines] = useState([])
  const [selectedEdge, setSelectedEdge] = useState(null)
  const [oosInput, setOosInput] = useState('')
  const [applyingOos, setApplyingOos] = useState(false)
  const [warning, setWarning] = useState(null)

  // Persist system choice in URL
  const changeSystem = (newSystem) => {
    setSystem(newSystem)
    setSearchParams({ system: newSystem })
  }

  useEffect(() => {
    if (!user) {
      navigate('/login')
      return
    }
    api.setToken(token)
    loadGrid()
  }, [user, token, system])

  const loadGrid = async () => {
    setLoading(true)
    setError(null)
    try {
      const [gridResp, baselineResp] = await Promise.all([
        api.getGrid(system),
        api.getBaseline(system)
      ])
      
      setGridData(gridResp)
      setPowerFlow(baselineResp.power_flow)
      
      // Extract open lines
      const openLineIds = gridResp.branches
        .filter(b => !b.data.in_service)
        .map(b => b.data.lineId)
      setOpenLines(openLineIds)
      
      // Set nodes
      setNodes(gridResp.nodes.map(n => ({
        ...n,
        type: 'busNode',
      })))
      
      // Set edges with dynamic styling
      setEdges(gridResp.branches.map(b => ({
        id: b.id,
        source: b.source,
        target: b.target,
        animated: !b.data.in_service,
        style: {
          stroke: b.data.in_service ? '#22c55e' : '#ef4444',
          strokeWidth: 2.5,
        },
        markerEnd: {
          type: MarkerType.Arrow,
          color: b.data.in_service ? '#22c55e' : '#ef4444',
        },
        data: b.data,
      })))
    } catch (err) {
      setError(err.message)
    } finally {
      setLoading(false)
    }
  }

  const onEdgeClick = useCallback((event, edge) => {
    setSelectedEdge(edge)
  }, [])

  const toggleSwitch = async () => {
    if (!selectedEdge) return
    
    setSimulating(true)
    setError(null)
    setSuccess(null)
    setWarning(null)
    
    try {
      const lineId = selectedEdge.data.lineId
      const result = await api.toggleSwitch({
        system,
        line_id: lineId,
        current_open_lines: openLines
      })
      
      if (result.is_valid) {
        setOpenLines(result.new_open_lines)
        if (result.power_flow) {
          setPowerFlow(result.power_flow)
        }
        
        // Update edges
        setEdges(eds => eds.map(e => {
          if (e.id === selectedEdge.id) {
            const isOpen = result.new_open_lines.includes(e.data.lineId)
            return {
              ...e,
              animated: isOpen,
              style: {
                ...e.style,
                stroke: isOpen ? '#ef4444' : '#22c55e',
              },
              markerEnd: {
                type: MarkerType.Arrow,
                color: isOpen ? '#ef4444' : '#22c55e',
              },
              data: {
                ...e.data,
                in_service: !isOpen
              }
            }
          }
          return e
        }))
        
        setSuccess(`Switch ${result.action} successfully. Loss: ${result.power_flow?.total_loss_kw?.toFixed(2)} kW`)
        if (result.warnings && result.warnings.length > 0) {
          setWarning(result.warnings.join(' '))
        }
        setSelectedEdge(null)
      } else {
        setError(result.error || 'Invalid configuration')
      }
    } catch (err) {
      setError(err.message)
    } finally {
      setSimulating(false)
    }
  }

  const runOptimization = async () => {
    setOptimizing(true)
    setError(null)
    setSuccess(null)
    setWarning(null)
    
    try {
      const result = await api.optimize({
        system,
        method: 'hybrid',
        open_lines: openLines.length > 0 ? openLines : undefined,
      })
      
      if (result.optimized) {
        const newOpenLines = result.optimized.open_lines || []
        setOpenLines(newOpenLines)
        setPowerFlow(result.optimized)
        
        // Update edges
        setEdges(eds => eds.map(e => {
          const isOpen = newOpenLines.includes(e.data.lineId)
          return {
            ...e,
            animated: isOpen,
            style: {
              ...e.style,
              stroke: isOpen ? '#ef4444' : '#22c55e',
            },
            markerEnd: {
              type: MarkerType.Arrow,
              color: isOpen ? '#ef4444' : '#22c55e',
            },
            data: {
              ...e.data,
              in_service: !isOpen
            }
          }
        }))
        
        setSuccess(`Optimized! Loss reduced to ${result.optimized.total_loss_kw.toFixed(2)} kW (${result.impact.loss_reduction_pct.toFixed(1)}% reduction)`)
      }
    } catch (err) {
      setError(err.message)
    } finally {
      setOptimizing(false)
    }
  }

  const downloadReport = async () => {
    setDownloadingReport(true)
    try {
      const html = await api.generateReport({ system, method: 'classical', include_optimization: true })
      const blob = new Blob([html], { type: 'text/html' })
      const url = URL.createObjectURL(blob)
      const a = document.createElement('a')
      a.href = url
      a.download = 'optiq_report.html'
      a.click()
      URL.revokeObjectURL(url)
    } catch (err) {
      setError(err.message)
    } finally {
      setDownloadingReport(false)
    }
  }

  const addOosLine = () => {
    const lineId = parseInt(oosInput.trim(), 10)
    if (isNaN(lineId)) return
    if (openLines.includes(lineId)) {
      setError(`Line ${lineId} is already out of service`)
      return
    }
    setOosInput('')
    // Don't apply yet - just add to staging list
    setOpenLines(prev => [...prev, lineId])
  }

  const removeOosLine = (lineId) => {
    setOpenLines(prev => prev.filter(l => l !== lineId))
  }

  const applyOutOfService = async () => {
    setApplyingOos(true)
    setError(null)
    setSuccess(null)
    setWarning(null)

    try {
      const result = await api.setOutOfServiceLines({
        system,
        out_of_service_lines: openLines,
      })

      if (result.valid) {
        // Sync open lines from response (may differ due to auto-repair)
        const actualOos = result.open_lines || openLines
        setOpenLines(actualOos)
        setPowerFlow(result.power_flow)

        // Update edge visuals
        setEdges(eds => eds.map(e => {
          const isOpen = actualOos.includes(e.data.lineId)
          return {
            ...e,
            animated: isOpen,
            style: {
              ...e.style,
              stroke: isOpen ? '#ef4444' : '#22c55e',
            },
            markerEnd: {
              type: MarkerType.Arrow,
              color: isOpen ? '#ef4444' : '#22c55e',
            },
            data: {
              ...e.data,
              in_service: !isOpen,
            },
          }
        }))

        let msg = `Applied ${openLines.length} out-of-service lines. Loss: ${result.power_flow?.total_loss_kw?.toFixed(2)} kW`
        if (result.warnings && result.warnings.length > 0) {
          setWarning(result.warnings.join(' '))
        }
        setSuccess(msg)
      } else {
        setError(result.error || 'Invalid configuration')
      }
    } catch (err) {
      setError(err.message)
    } finally {
      setApplyingOos(false)
    }
  }

  return (
    <div className="h-[calc(100vh-64px)] flex flex-col">
      {/* Controls Bar */}
      <div className="bg-white border-b border-dark-200 px-4 py-3">
        <div className="max-w-7xl mx-auto flex flex-wrap items-center justify-between gap-4">
          <div className="flex items-center gap-4">
            <h1 className="font-display text-xl font-bold text-dark-800">Grid View</h1>
            <select
              value={system}
              onChange={(e) => changeSystem(e.target.value)}
              className="input w-auto text-sm py-2"
            >
              <option value="case33bw">IEEE 33-Bus</option>
              <option value="case118">IEEE 118-Bus</option>
            </select>
          </div>
          
          <div className="flex items-center gap-3">
            <button
              onClick={loadGrid}
              disabled={loading}
              className="btn btn-outline text-sm py-2"
            >
              <RefreshCw className={`w-4 h-4 mr-2 ${loading ? 'animate-spin' : ''}`} />
              Refresh
            </button>
            
            <button
              onClick={runOptimization}
              disabled={optimizing || loading}
              className="btn btn-primary text-sm py-2"
            >
              {optimizing ? (
                <>
                  <RefreshCw className="w-4 h-4 mr-2 animate-spin" />
                  Optimizing...
                </>
              ) : (
                <>
                  <Zap className="w-4 h-4 mr-2" />
                  Optimize
                </>
              )}
            </button>
            
            <button
              onClick={downloadReport}
              disabled={downloadingReport}
              className="btn btn-secondary text-sm py-2"
            >
              {downloadingReport ? (
                <Loader2 className="w-4 h-4 mr-2 animate-spin" />
              ) : (
                <Download className="w-4 h-4 mr-2" />
              )}
              {downloadingReport ? 'Preparing...' : 'Report'}
            </button>
          </div>
        </div>
      </div>

      {/* Alerts */}
      {error && (
        <div className="bg-red-50 border-b border-red-200 px-4 py-2 text-red-600 flex items-center gap-2">
          <AlertCircle className="w-4 h-4" />
          {error}
          <button onClick={() => setError(null)} className="ml-auto text-red-400 hover:text-red-600">×</button>
        </div>
      )}
      {warning && (
        <div className="bg-amber-50 border-b border-amber-200 px-4 py-2 text-amber-700 flex items-center gap-2">
          <AlertCircle className="w-4 h-4" />
          {warning}
          <button onClick={() => setWarning(null)} className="ml-auto text-amber-400 hover:text-amber-600">×</button>
        </div>
      )}
      {success && (
        <div className="bg-green-50 border-b border-green-200 px-4 py-2 text-green-600 flex items-center gap-2">
          <CheckCircle2 className="w-4 h-4" />
          {success}
          <button onClick={() => setSuccess(null)} className="ml-auto text-green-400 hover:text-green-600">×</button>
        </div>
      )}

      {/* Main Content */}
      <div className="flex-1 flex">
        {/* React Flow */}
        <div className="flex-1 relative">
          {loading ? (
            <div className="absolute inset-0 flex items-center justify-center bg-dark-50">
              <RefreshCw className="w-8 h-8 text-primary-500 animate-spin" />
            </div>
          ) : (
            <ReactFlow
              nodes={nodes}
              edges={edges}
              onNodesChange={onNodesChange}
              onEdgesChange={onEdgesChange}
              onEdgeClick={onEdgeClick}
              nodeTypes={nodeTypes}
              fitView
              attributionPosition="bottom-left"
            >
              <Background color="#e2e8f0" gap={20} />
              <Controls />
              <MiniMap 
                nodeColor={(node) => node.data?.is_slack ? '#0ea5e9' : '#94a3b8'}
                maskColor="rgba(0, 0, 0, 0.1)"
              />
            </ReactFlow>
          )}
        </div>

        {/* Side Panel */}
        <div className="w-80 bg-white border-l border-dark-200 overflow-y-auto">
          {/* Power Flow Info */}
          {powerFlow && (
            <div className="p-4 border-b border-dark-200">
              <h3 className="font-semibold text-dark-800 mb-3">Power Flow Results</h3>
              <div className="space-y-2 text-sm">
                <div className="flex justify-between">
                  <span className="text-dark-500">Total Loss:</span>
                  <span className="font-medium">{powerFlow.total_loss_kw?.toFixed(2)} kW</span>
                </div>
                <div className="flex justify-between">
                  <span className="text-dark-500">Loss %:</span>
                  <span className="font-medium">{powerFlow.loss_pct?.toFixed(2)}%</span>
                </div>
                <div className="flex justify-between">
                  <span className="text-dark-500">Min Voltage:</span>
                  <span className={`font-medium ${powerFlow.min_voltage_pu < 0.95 ? 'text-red-600' : 'text-green-600'}`}>
                    {powerFlow.min_voltage_pu?.toFixed(4)} pu
                  </span>
                </div>
                <div className="flex justify-between">
                  <span className="text-dark-500">Violations:</span>
                  <span className={`font-medium ${powerFlow.voltage_violations > 0 ? 'text-red-600' : 'text-green-600'}`}>
                    {powerFlow.voltage_violations}
                  </span>
                </div>
              </div>
            </div>
          )}

          {/* Selected Edge */}
          {selectedEdge && (
            <div className="p-4 border-b border-dark-200 bg-primary-50">
              <h3 className="font-semibold text-dark-800 mb-3 flex items-center gap-2">
                <Info className="w-4 h-4 text-primary-500" />
                Selected Line
              </h3>
              <div className="space-y-2 text-sm mb-4">
                <div className="flex justify-between">
                  <span className="text-dark-500">Line ID:</span>
                  <span className="font-medium">{selectedEdge.data.lineId}</span>
                </div>
                <div className="flex justify-between">
                  <span className="text-dark-500">From Bus:</span>
                  <span className="font-medium">{selectedEdge.data.from_bus}</span>
                </div>
                <div className="flex justify-between">
                  <span className="text-dark-500">To Bus:</span>
                  <span className="font-medium">{selectedEdge.data.to_bus}</span>
                </div>
                <div className="flex justify-between">
                  <span className="text-dark-500">Status:</span>
                  <span className={`font-medium ${selectedEdge.data.in_service ? 'text-green-600' : 'text-red-600'}`}>
                    {selectedEdge.data.in_service ? 'CLOSED' : 'OPEN'}
                  </span>
                </div>
              </div>
              
              <button
                onClick={toggleSwitch}
                disabled={simulating}
                className={`btn w-full ${selectedEdge.data.in_service ? 'btn-secondary' : 'btn-accent'}`}
              >
                {simulating ? (
                  <RefreshCw className="w-4 h-4 mr-2 animate-spin" />
                ) : selectedEdge.data.in_service ? (
                  <ToggleRight className="w-4 h-4 mr-2" />
                ) : (
                  <ToggleLeft className="w-4 h-4 mr-2" />
                )}
                {selectedEdge.data.in_service ? 'Open Switch' : 'Close Switch'}
              </button>
            </div>
          )}

          {/* Out of Service Lines Management */}
          <div className="p-4 border-b border-dark-200">
            <h3 className="font-semibold text-dark-800 mb-3 flex items-center gap-2">
              <Power className="w-4 h-4 text-red-500" />
              Out of Service Lines ({openLines.length})
            </h3>
            
            {/* Add line input */}
            <div className="flex gap-2 mb-3">
              <input
                type="number"
                min="0"
                value={oosInput}
                onChange={(e) => setOosInput(e.target.value)}
                onKeyDown={(e) => e.key === 'Enter' && addOosLine()}
                placeholder="Line ID"
                className="input text-sm py-1.5 flex-1"
              />
              <button
                onClick={addOosLine}
                className="btn btn-outline text-sm py-1.5 px-2"
                title="Add line"
              >
                <Plus className="w-4 h-4" />
              </button>
            </div>

            {/* Current out-of-service lines */}
            {openLines.length > 0 ? (
              <div className="flex flex-wrap gap-1.5 mb-3">
                {openLines.map(lineId => (
                  <span
                    key={lineId}
                    className="inline-flex items-center gap-1 px-2 py-0.5 bg-red-100 text-red-700 rounded text-sm font-medium group"
                  >
                    Line {lineId}
                    <button
                      onClick={() => removeOosLine(lineId)}
                      className="opacity-60 hover:opacity-100"
                      title={`Remove line ${lineId}`}
                    >
                      <X className="w-3 h-3" />
                    </button>
                  </span>
                ))}
              </div>
            ) : (
              <p className="text-sm text-dark-500 mb-3">No out-of-service lines</p>
            )}

            {/* Apply button */}
            <button
              onClick={applyOutOfService}
              disabled={applyingOos || loading}
              className="btn btn-secondary w-full text-sm py-2"
            >
              {applyingOos ? (
                <>
                  <Loader2 className="w-4 h-4 mr-2 animate-spin" />
                  Applying...
                </>
              ) : (
                <>
                  <Play className="w-4 h-4 mr-2" />
                  Apply &amp; Run Power Flow
                </>
              )}
            </button>
          </div>

          {/* Legend */}
          <div className="p-4 border-t border-dark-200">
            <h3 className="font-semibold text-dark-800 mb-3">Legend</h3>
            <div className="space-y-2 text-sm">
              <div className="flex items-center gap-2">
                <div className="w-6 h-0.5 bg-green-500"></div>
                <span className="text-dark-500">Closed (In Service)</span>
              </div>
              <div className="flex items-center gap-2">
                <div className="w-6 h-0.5 bg-red-500"></div>
                <span className="text-dark-500">Open (Out of Service)</span>
              </div>
              <div className="flex items-center gap-2">
                <div className="w-4 h-4 rounded bg-gradient-to-br from-primary-500 to-primary-600"></div>
                <span className="text-dark-500">Slack Bus (Generator)</span>
              </div>
              <div className="flex items-center gap-2">
                <div className="w-4 h-4 rounded bg-white border-2 border-dark-200"></div>
                <span className="text-dark-500">Load Bus</span>
              </div>
            </div>
          </div>
        </div>
      </div>
    </div>
  )
}