//! Comprehensive tests for append functionality

#[cfg(test)]
mod append_tests {
    use ndarray::Array2;
    use taco_format::{
        AtomMetadata, CompressionSettings, Frame, FrameData, Header, PrecisionMode, Reader,
        SimulationMetadata, Writer,
    };
    use tempfile::tempdir;

    #[test]
    fn test_append_maintains_delta_chain() {
        // This test verifies that the delta encoding chain is properly maintained
        // when appending frames
        let dir = tempdir().expect("Failed to create temp dir");
        let file_path = dir.path().join("test_delta_chain.taco");

        let sim_metadata = SimulationMetadata::default();
        let atom_metadata = AtomMetadata::default();
        let compression_settings = CompressionSettings::default();

        // Create initial file with 3 frames
        let mut writer = Writer::create(
            &file_path,
            2,
            0.001,
            sim_metadata,
            atom_metadata,
            compression_settings,
        )
        .expect("Failed to create writer");

        let mut positions = Array2::<f32>::zeros((2, 3));
        for i in 0..3 {
            for j in 0..2 {
                for k in 0..3 {
                    // Create a slowly changing trajectory
                    positions[[j, k]] = (i as f32) * 0.01 + (j as f32) + (k as f32) * 0.1;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Append 4 more frames - these should use frame 2 as reference
        let mut writer = Writer::append(&file_path).expect("Failed to append");
        for i in 3..7 {
            for j in 0..2 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) * 0.01 + (j as f32) + (k as f32) * 0.1;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Verify all frames can be read correctly
        let mut reader = Reader::open(&file_path).expect("Failed to open reader");
        assert_eq!(reader.num_frames(), 7);

        for i in 0..7 {
            let frame = reader.read_frame(i).expect("Failed to read frame");
            if let Some(pos) = &frame.data.positions {
                for j in 0..2 {
                    for k in 0..3 {
                        let expected = (i as f32) * 0.01 + (j as f32) + (k as f32) * 0.1;
                        let actual = pos[[j, k]];
                        assert!(
                            (actual - expected).abs() < 0.001,
                            "Frame {}, atom {}, coord {}: expected {}, got {}, diff {}",
                            i,
                            j,
                            k,
                            expected,
                            actual,
                            (actual - expected).abs()
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn test_append_respects_full_frame_interval() {
        let dir = tempdir().expect("Failed to create temp dir");
        let file_path = dir.path().join("test_full_interval.taco");

        let sim_metadata = SimulationMetadata::default();
        let atom_metadata = AtomMetadata::default();
        let compression_settings = CompressionSettings::default();

        // Create header with full_frame_interval = 5
        let mut header = Header::new(2, 0.001, sim_metadata, atom_metadata, compression_settings);
        header.full_frame_interval = 5;

        let mut writer =
            Writer::create_with_header(&file_path, header).expect("Failed to create writer");

        // Write 6 frames (0-5, where 0 and 5 should be full frames)
        let mut positions = Array2::<f32>::zeros((2, 3));
        for i in 0..6 {
            for j in 0..2 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) + (j as f32) * 0.1;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Append 7 more frames (6-12, where 10 should be a full frame)
        let mut writer = Writer::append(&file_path).expect("Failed to append");
        for i in 6..13 {
            for j in 0..2 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) + (j as f32) * 0.1;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Verify frame types
        let mut reader = Reader::open(&file_path).expect("Failed to open reader");
        assert_eq!(reader.num_frames(), 13);

        for i in 0..13 {
            let frame = reader.read_frame(i).expect("Failed to read frame");
            let expected_full = i % 5 == 0;
            assert_eq!(
                frame.is_full_frame,
                expected_full,
                "Frame {} should be {} but was {}",
                i,
                if expected_full { "full" } else { "delta" },
                if frame.is_full_frame { "full" } else { "delta" }
            );

            // Verify data is correct
            if let Some(pos) = &frame.data.positions {
                for j in 0..2 {
                    for k in 0..3 {
                        let expected = (i as f32) + (j as f32) * 0.1;
                        let actual = pos[[j, k]];
                        assert!(
                            (actual - expected).abs() < 0.001,
                            "Frame {} position mismatch",
                            i
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn test_append_preserves_metadata() {
        let dir = tempdir().expect("Failed to create temp dir");
        let file_path = dir.path().join("test_metadata.taco");

        let sim_metadata = SimulationMetadata {
            name: Some("Test Simulation".to_string()),
            temperature: Some(300.0),
            pressure: Some(1.0),
            ensemble: Some("NPT".to_string()),
            ..Default::default()
        };

        let atom_metadata = AtomMetadata {
            names: Some(vec!["H".to_string(), "O".to_string()]),
            elements: Some(vec!["H".to_string(), "O".to_string()]),
            masses: Some(vec![1.008, 16.0]),
            ..Default::default()
        };

        let compression_settings = CompressionSettings {
            precision: PrecisionMode::Half,
            ..Default::default()
        };

        // Create initial file
        let mut writer = Writer::create(
            &file_path,
            2,
            0.002,
            sim_metadata.clone(),
            atom_metadata.clone(),
            compression_settings.clone(),
        )
        .expect("Failed to create writer");

        let positions = Array2::<f32>::zeros((2, 3));
        let frame = Frame::new(0, 0.0, FrameData::new(positions));
        writer.write_frame(frame).expect("Failed to write frame");
        writer.finish().expect("Failed to finish");

        // Append to file
        let mut writer = Writer::append(&file_path).expect("Failed to append");
        let positions = Array2::<f32>::ones((2, 3));
        let frame = Frame::new(1, 0.002, FrameData::new(positions));
        writer.write_frame(frame).expect("Failed to write frame");
        writer.finish().expect("Failed to finish");

        // Verify metadata is preserved
        let reader = Reader::open(&file_path).expect("Failed to open reader");
        let header = reader.header();

        assert_eq!(
            header.simulation_metadata.name,
            Some("Test Simulation".to_string())
        );
        assert_eq!(header.simulation_metadata.temperature, Some(300.0));
        assert_eq!(header.simulation_metadata.pressure, Some(1.0));
        assert_eq!(header.simulation_metadata.ensemble, Some("NPT".to_string()));

        assert_eq!(
            header.atom_metadata.names,
            Some(vec!["H".to_string(), "O".to_string()])
        );
        assert_eq!(
            header.atom_metadata.elements,
            Some(vec!["H".to_string(), "O".to_string()])
        );
        assert_eq!(header.atom_metadata.masses, Some(vec![1.008, 16.0]));

        assert_eq!(header.num_atoms, 2);
        assert_eq!(header.time_step, 0.002);
        assert_eq!(header.num_frames, 2);
    }

    #[test]
    fn test_append_with_complex_data() {
        let dir = tempdir().expect("Failed to create temp dir");
        let file_path = dir.path().join("test_complex.taco");

        let sim_metadata = SimulationMetadata::default();
        let atom_metadata = AtomMetadata::default();
        let compression_settings = CompressionSettings::default();

        // Create initial file with positions, velocities, and forces
        let mut writer = Writer::create(
            &file_path,
            3,
            0.001,
            sim_metadata,
            atom_metadata,
            compression_settings,
        )
        .expect("Failed to create writer");

        for i in 0..3 {
            let mut positions = Array2::<f32>::zeros((3, 3));
            let mut velocities = Array2::<f32>::zeros((3, 3));
            let mut forces = Array2::<f32>::zeros((3, 3));

            for j in 0..3 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) * 0.1 + (j as f32) + (k as f32) * 0.01;
                    velocities[[j, k]] = (i as f32) * 0.5 + (j as f32) * 0.05;
                    forces[[j, k]] = (i as f32) * 2.0 + (j as f32) * 0.2;
                }
            }

            let frame_data = FrameData {
                positions: Some(positions),
                velocities: Some(velocities),
                forces: Some(forces),
                stress: None,
                extra: std::collections::HashMap::new(),
            };

            let frame = Frame::new(i, i as f64 * 0.001, frame_data)
                .with_box_dimensions([10.0, 10.0, 10.0], [90.0, 90.0, 90.0])
                .with_thermo(300.0 + i as f32, 1.0 + i as f32 * 0.1);

            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Append more complex frames
        let mut writer = Writer::append(&file_path).expect("Failed to append");
        for i in 3..6 {
            let mut positions = Array2::<f32>::zeros((3, 3));
            let mut velocities = Array2::<f32>::zeros((3, 3));
            let mut forces = Array2::<f32>::zeros((3, 3));

            for j in 0..3 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) * 0.1 + (j as f32) + (k as f32) * 0.01;
                    velocities[[j, k]] = (i as f32) * 0.5 + (j as f32) * 0.05;
                    forces[[j, k]] = (i as f32) * 2.0 + (j as f32) * 0.2;
                }
            }

            let frame_data = FrameData {
                positions: Some(positions),
                velocities: Some(velocities),
                forces: Some(forces),
                stress: None,
                extra: std::collections::HashMap::new(),
            };

            let frame = Frame::new(i, i as f64 * 0.001, frame_data)
                .with_box_dimensions([10.0, 10.0, 10.0], [90.0, 90.0, 90.0])
                .with_thermo(300.0 + i as f32, 1.0 + i as f32 * 0.1);

            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Verify all data
        let mut reader = Reader::open(&file_path).expect("Failed to open reader");
        assert_eq!(reader.num_frames(), 6);

        for i in 0..6 {
            let frame = reader.read_frame(i).expect("Failed to read frame");

            // Check positions
            if let Some(pos) = &frame.data.positions {
                for j in 0..3 {
                    for k in 0..3 {
                        let expected = (i as f32) * 0.1 + (j as f32) + (k as f32) * 0.01;
                        let actual = pos[[j, k]];
                        assert!(
                            (actual - expected).abs() < 0.001,
                            "Frame {} position mismatch",
                            i
                        );
                    }
                }
            }

            // Check velocities
            if let Some(vel) = &frame.data.velocities {
                for j in 0..3 {
                    for k in 0..3 {
                        let expected = (i as f32) * 0.5 + (j as f32) * 0.05;
                        let actual = vel[[j, k]];
                        assert!(
                            (actual - expected).abs() < 0.01,
                            "Frame {} velocity mismatch",
                            i
                        );
                    }
                }
            }

            // Check forces
            if let Some(force) = &frame.data.forces {
                for j in 0..3 {
                    for k in 0..3 {
                        let expected = (i as f32) * 2.0 + (j as f32) * 0.2;
                        let actual = force[[j, k]];
                        assert!(
                            (actual - expected).abs() < 0.01,
                            "Frame {} force mismatch",
                            i
                        );
                    }
                }
            }

            // Check box dimensions
            if let Some(box_dims) = &frame.box_dimensions {
                for j in 0..3 {
                    assert!((box_dims[j] - 10.0).abs() < 0.001);
                }
            }

            // Check thermodynamic properties
            if let Some(temp) = frame.temperature {
                assert!((temp - (300.0 + i as f32)).abs() < 0.1);
            }
            if let Some(pressure) = frame.pressure {
                assert!((pressure - (1.0 + i as f32 * 0.1)).abs() < 0.01);
            }
        }
    }

    #[test]
    fn test_append_large_trajectory() {
        let dir = tempdir().expect("Failed to create temp dir");
        let file_path = dir.path().join("test_large.taco");

        let sim_metadata = SimulationMetadata::default();
        let atom_metadata = AtomMetadata::default();
        let compression_settings = CompressionSettings::default();

        // Create initial file with 50 frames
        let mut writer = Writer::create(
            &file_path,
            10,
            0.001,
            sim_metadata,
            atom_metadata,
            compression_settings,
        )
        .expect("Failed to create writer");

        let mut positions = Array2::<f32>::zeros((10, 3));
        for i in 0..50 {
            for j in 0..10 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) * 0.01 + (j as f32) * 0.1 + (k as f32) * 0.001;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Append 50 more frames
        let mut writer = Writer::append(&file_path).expect("Failed to append");
        for i in 50..100 {
            for j in 0..10 {
                for k in 0..3 {
                    positions[[j, k]] = (i as f32) * 0.01 + (j as f32) * 0.1 + (k as f32) * 0.001;
                }
            }
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
        }
        writer.finish().expect("Failed to finish");

        // Verify total frame count and random access
        let mut reader = Reader::open(&file_path).expect("Failed to open reader");
        assert_eq!(reader.num_frames(), 100);

        // Test random access to various frames
        for &frame_idx in &[0, 25, 49, 50, 75, 99] {
            let frame = reader
                .read_frame(frame_idx)
                .expect(&format!("Failed to read frame {}", frame_idx));
            assert_eq!(frame.frame_number, frame_idx);

            if let Some(pos) = &frame.data.positions {
                for j in 0..10 {
                    for k in 0..3 {
                        let expected =
                            (frame_idx as f32) * 0.01 + (j as f32) * 0.1 + (k as f32) * 0.001;
                        let actual = pos[[j, k]];
                        assert!(
                            (actual - expected).abs() < 0.001,
                            "Frame {} mismatch",
                            frame_idx
                        );
                    }
                }
            }
        }
    }

    #[test]
    fn test_multiple_sequential_appends() {
        let dir = tempdir().expect("Failed to create temp dir");
        let file_path = dir.path().join("test_sequential_appends.taco");

        let sim_metadata = SimulationMetadata::default();
        let atom_metadata = AtomMetadata::default();
        let compression_settings = CompressionSettings::default();

        // Create initial file with 1 frame
        let mut writer = Writer::create(
            &file_path,
            2,
            0.001,
            sim_metadata,
            atom_metadata,
            compression_settings,
        )
        .expect("Failed to create writer");

        let mut positions = Array2::<f32>::zeros((2, 3));
        positions[[0, 0]] = 0.0;
        let frame = Frame::new(0, 0.0, FrameData::new(positions.clone()));
        writer.write_frame(frame).expect("Failed to write frame");
        writer.finish().expect("Failed to finish");

        // Append 10 times, each time adding 1 frame
        for i in 1..11 {
            let mut writer = Writer::append(&file_path).expect("Failed to append");
            positions[[0, 0]] = i as f32;
            let frame = Frame::new(i, i as f64 * 0.001, FrameData::new(positions.clone()));
            writer.write_frame(frame).expect("Failed to write frame");
            writer.finish().expect("Failed to finish");
        }

        // Verify we have 11 frames total
        let mut reader = Reader::open(&file_path).expect("Failed to open reader");
        assert_eq!(reader.num_frames(), 11);

        for i in 0..11 {
            let frame = reader.read_frame(i).expect("Failed to read frame");
            assert_eq!(frame.frame_number, i);
            if let Some(pos) = &frame.data.positions {
                assert!((pos[[0, 0]] - i as f32).abs() < 0.001);
            }
        }
    }
}
