import React, { useState, useEffect } from 'react';
import axios from 'axios';
import { Container, Navbar, Card, ProgressBar, Alert, Button, Form, Row, Col, ListGroup, Badge } from 'react-bootstrap';

const API_URL = process.env.NODE_ENV === 'development' ? 'http://localhost:8000' : '';
const WS_URL = process.env.NODE_ENV === 'development' 
    ? 'ws://localhost:8000' 
    : window.location.protocol.replace('http', 'ws') + '//' + window.location.host;

// --- Individual Job Item Component ---
const JobItem = ({ job, API_URL, onJobDelete }) => {
  const getVariant = (status) => {
    switch (status) {
      case 'PENDING': return 'info';
      case 'PROCESSING': return 'primary';
      case 'COMPLETE': return 'success';
      case 'FAILED': return 'danger';
      default: return 'secondary';
    }
  };

  const handleDelete = async () => {
    if (window.confirm(`Are you sure you want to delete job ${job.filename} (${job.id})?`)) {
      try {
        await axios.delete(`${API_URL}/api/jobs/${job.id}`);
        onJobDelete(job.id);
      } catch (error) {
        console.error("Error deleting job:", error);
        alert("Failed to delete job.");
      }
    }
  };

  return (
    <ListGroup.Item>
      <div className="d-flex justify-content-between align-items-center">
        <div>
          <strong>{job.filename}</strong> (ID: {job.id.substring(0, 8) + '...'})
          <br />
          <small className="text-muted">Status: <Badge bg={getVariant(job.status)}>{job.status}</Badge></small>
        </div>
        <div>
          {job.status === 'COMPLETE' && job.download_url && (
            <Button variant="success" href={`${API_URL}${job.download_url}`} size="sm" className="me-2">
              Download
            </Button>
          )}
          <Button variant="outline-danger" size="sm" onClick={handleDelete}>
            Delete
          </Button>
        </div>
      </div>

      {/* Progress Bar for Processing state */}
      {job.status === 'PROCESSING' && <ProgressBar now={job.progress} label={`${job.progress}%`} animated className="mt-2" />}
      
      {/* Informational Message for non-failed states */}
      {job.status !== 'FAILED' && job.message && (
        <small className="d-block mt-1 text-muted">{job.message}</small>
      )}

      {/* Error Alert for FAILED state */}
      {job.status === 'FAILED' && job.message && (
        <Alert variant="danger" className="mt-2 mb-0" style={{ fontSize: '0.875em', padding: '0.5rem 1rem' }}>
          <strong>Error:</strong> {job.message}
        </Alert>
      )}
    </ListGroup.Item>
  );
};

// --- Upload Component (modified) ---
const UploadComponent = ({ setJobs }) => {
  const [file, setFile] = useState(null);
  const [numLayers, setNumLayers] = useState(20);
  const [numPoints, setNumPoints] = useState(30);
  const [uploadProgress, setUploadProgress] = useState(0); // For file upload progress
  const [uploadStatusMessage, setUploadStatusMessage] = useState('Select a file to begin.');
  const [error, setError] = useState(null);
  const [isUploading, setIsUploading] = useState(false);

  const handleFileChange = (event) => {
    setFile(event.target.files[0]);
    setUploadStatusMessage(event.target.files[0] ? event.target.files[0].name : 'Select a file to begin.');
    setUploadProgress(0);
    setError(null);
  };

  const handleUpload = async () => {
    if (!file) return;

    setIsUploading(true);
    setUploadProgress(0);
    setError(null);
    setUploadStatusMessage('Uploading file...');

    const formData = new FormData();
    formData.append('file', file);
    formData.append('num_layers', numLayers);
    formData.append('num_points_per_layer', numPoints);

    try {
      const uploadResponse = await axios.post(`${API_URL}/upload/`, formData, {
        onUploadProgress: (progressEvent) => {
          const percentCompleted = Math.round((progressEvent.loaded * 100) / progressEvent.total);
          setUploadProgress(percentCompleted);
          setUploadStatusMessage(`Uploading... ${percentCompleted}%`);
        },
      });

      const { job_id, filename, status } = uploadResponse.data;
      // Add the new job to the global state immediately
      setJobs(prevJobs => [{
        id: job_id,
        filename: filename,
        status: status,
        progress: 0,
        message: `File '${filename}' uploaded, awaiting processing.`, 
        num_layers: numLayers,
        num_points_per_layer: numPoints,
        timestamp: new Date().toISOString() // Use ISO string for consistency
      }, ...prevJobs]);

      setUploadStatusMessage(`Job ${job_id.substring(0, 8)}... created. Waiting for WebSocket updates.`);
      setUploadProgress(0); // Reset upload progress for next file
      setFile(null); // Clear the file input

      // No longer open WebSocket here. WebSocket will be opened by JobList for ongoing jobs.

    } catch (err) {
      console.error(err);
      setError(err.response?.data?.detail || 'File upload failed. Please try again.');
    } finally {
      setIsUploading(false);
    }
  };

  return (
    <Card className="mb-4">
      <Card.Body>
        <Card.Title>Upload Mesh File</Card.Title>
        <Card.Text>
          Select a <code>.obj</code>, <code>.stl</code>, or <code>.3mf</code> file to process. The tool will generate a DXF file containing the layered profiles of your model.
        </Card.Text>
        <Form>
          <Form.Group controlId="formFile" className="mb-3">
            <Form.Label>Select Mesh File</Form.Label>
            <Form.Control type="file" onChange={handleFileChange} accept=".obj,.stl,.3mf" />
          </Form.Group>

          <Row className="mb-3">
            <Form.Group as={Col} controlId="formNumLayers">
              <Form.Label>Number of Layers</Form.Label>
              <Form.Control 
                type="number" 
                value={numLayers} 
                onChange={(e) => setNumLayers(parseInt(e.target.value))} 
                min="2" 
                max="100" 
              />
            </Form.Group>
            <Form.Group as={Col} controlId="formNumPoints">
              <Form.Label>Points per Layer</Form.Label>
              <Form.Control 
                type="number" 
                value={numPoints} 
                onChange={(e) => setNumPoints(parseInt(e.target.value))} 
                min="3" 
                max="200" 
              />
            </Form.Group>
          </Row>
          
          {file && (
            <Button variant="primary" onClick={handleUpload} className="mb-3" disabled={isUploading}>
              {isUploading ? 'Processing...' : 'Start Processing'}
            </Button>
          )}
        </Form>

        <hr />

        <h6>Upload Status</h6>
        <p className="text-muted">{uploadStatusMessage}</p>
        {uploadProgress > 0 && uploadProgress < 100 && (
            <ProgressBar now={uploadProgress} label={`${uploadProgress}%`} animated />
        )}

        {error && (
          <Alert variant="danger" className="mt-3">
            {error}
          </Alert>
        )}
      </Card.Body>
    </Card>
  );
};


// --- Main App Component (with WebSocket logic) ---
function App() {
  const [jobs, setJobs] = useState([]);
  const webSocketRef = React.useRef(null);

  // Fetch all existing jobs on component mount and set up polling
  useEffect(() => {
    const fetchJobs = async () => {
      try {
        const response = await axios.get(`${API_URL}/api/jobs`);
        // Sort jobs by timestamp on initial load
        const sortedJobs = response.data.sort((a, b) => new Date(b.timestamp) - new Date(a.timestamp));
        setJobs(sortedJobs);
      } catch (error) {
        console.error("Error fetching jobs:", error);
      }
    };
    
    fetchJobs(); // Initial fetch
    const interval = setInterval(fetchJobs, 5000); // Poll every 5 seconds

    return () => clearInterval(interval); // Cleanup polling on unmount
  }, []);

  // Effect to manage WebSocket connection for the primary job being tracked
  useEffect(() => {
    // A job is trackable if it's in a non-terminal state
    const jobToTrack = jobs.find(j => j.status === 'QUEUED' || j.status === 'PROCESSING');

    // If there is no job to track, or we are already connected to this job, do nothing.
    if (!jobToTrack || (webSocketRef.current && webSocketRef.current.jobId === jobToTrack.id)) {
      return;
    }

    // If there's an existing WebSocket for a different job, close it.
    if (webSocketRef.current) {
      webSocketRef.current.close();
    }

    // Create a new WebSocket connection
    const ws = new WebSocket(`${WS_URL}/ws/${jobToTrack.id}`);
    ws.jobId = jobToTrack.id; // Attach jobId for identification

    ws.onopen = () => {
      console.log(`WebSocket connected for job ${jobToTrack.id}`);
    };

    ws.onmessage = (event) => {
      const updatedJob = JSON.parse(event.data);
      setJobs(prevJobs => {
        // Create a map for quick lookup
        const jobsMap = new Map(prevJobs.map(j => [j.id, j]));
        // Update the job in the map
        jobsMap.set(updatedJob.id, { ...jobsMap.get(updatedJob.id), ...updatedJob });
        // Convert map back to array and re-sort
        return Array.from(jobsMap.values()).sort((a, b) => new Date(b.timestamp) - new Date(a.timestamp));
      });
    };
    
    // The 'onclose' event is expected when a job finishes.
    ws.onclose = (event) => {
      console.log(`WebSocket disconnected for job ${jobToTrack.id}. Code: ${event.code}, Reason: ${event.reason}`);
      // When the server closes the connection (e.g., job done), we nullify the ref.
      // The polling will handle fetching the final job state.
      if (webSocketRef.current && webSocketRef.current.jobId === jobToTrack.id) {
          webSocketRef.current = null;
      }
    };
    
    // The 'onerror' event is for unexpected connection failures.
    ws.onerror = (error) => {
      console.error(`WebSocket error for job ${jobToTrack.id}:`, error);
      // The onclose event will likely fire immediately after this.
    };

    webSocketRef.current = ws;

    // Cleanup function to close the WebSocket when the component unmounts
    // or when the job to track changes.
    return () => {
      if (ws && ws.readyState < 2) { // Only close if not already closing or closed
        console.log(`Cleaning up WebSocket for job ${ws.jobId}`);
        ws.close();
      }
    };
  }, [jobs]); // This effect re-evaluates whenever the list of jobs changes.

  const handleJobDelete = (jobId) => {
    setJobs(prevJobs => prevJobs.filter(job => job.id !== jobId));
  };

  const processingJobs = jobs.filter(j => j.status === 'PROCESSING' || j.status === 'QUEUED');
  const otherJobs = jobs.filter(j => j.status !== 'PROCESSING' && j.status !== 'QUEUED');
  
  return (
    <div>
      <Navbar bg="dark" variant="dark">
        <Container>
          <Navbar.Brand href="#home">DXF Curve Generator</Navbar.Brand>
        </Container>
      </Navbar>
      <Container className="mt-5">
        <UploadComponent setJobs={setJobs} />

        {processingJobs.length > 0 && (
          <Card className="mt-4">
            <Card.Body>
              <Card.Title>Currently Processing</Card.Title>
              <ListGroup variant="flush">
                {processingJobs.map(job => (
                   <JobItem key={job.id} job={job} API_URL={API_URL} onJobDelete={handleJobDelete} />
                ))}
              </ListGroup>
            </Card.Body>
          </Card>
        )}

        <Card className="mt-4">
          <Card.Body>
            <Card.Title>Job History</Card.Title>
            {otherJobs.length === 0 && processingJobs.length === 0 ? (
              <Card.Text>No jobs to show. Upload a file to begin.</Card.Text>
            ) : (
              <ListGroup variant="flush">
                {otherJobs.map(job => (
                  <JobItem key={job.id} job={job} API_URL={API_URL} onJobDelete={handleJobDelete} />
                ))}
              </ListGroup>
            )}
          </Card.Body>
        </Card>
      </Container>
    </div>
  );
}

export default App;