qred 0.1.2

use image::{ImageBuffer, Luma};

use crate::qr_spec::{ec_level_format_bits, get_alignment_coords, is_mask_bit_set};

// UTIL: Generate BCH(15, 5) format string
pub fn generate_format_string(ec_level: &str, mask_id: u8) -> u16 {
    let ec_bits = ec_level_format_bits(ec_level);
    let data = ((ec_bits as u16) << 3) | (mask_id as u16 & 0x07);

    let mut num = data << 10;
    let generator: u16 = 0x537;   // generator polynomial: x^10 + x^8 + x^5 + x^4 + x^2 + x + 1

    for i in (0..5).rev() {
        if (num & (1 << (i+10))) != 0 {
            num ^= generator << i;
        }
    }

    let formatstr = ((data << 10) | (num & 0x3FF)) ^ 0x5412;
    formatstr
}

// Matrix element structure
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum ModuleColor {
    Light,
    Dark,
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Module {
    pub color:          ModuleColor,
    pub is_reserved:    bool,
}

impl Module {
    pub fn default() -> Self {
        Self {
            color: ModuleColor::Light,
            is_reserved: false,
        }
    }
}

// Main matrix
pub struct Matrix {
    pub version:    usize,
    pub size:       usize,
    pub grid:       Vec<Vec<Module>>,
}

impl Matrix {
    pub fn new(version: usize) -> Self {
        let size = 17 + (4 * version);
        let grid = vec![vec![Module::default(); size]; size];
        let mut matrix = Self { version, size, grid };

        matrix.setup_static_patterns();
        matrix
    }

    // Set up static patterns
    fn setup_static_patterns(&mut self) {
        self.place_finders_and_separators();
        self.reserve_format_areas();
        self.place_dark_module();
        self.place_alignment_patterns();
        self.place_timing_patterns();
    }

    // Place data and EC codewords
    pub fn place_data_bits(&mut self, codewords: &[u8]) {
        let mut bits = Vec::with_capacity(codewords.len() * 8);
        for &byte in codewords {
            for i in (0..8).rev() {
                bits.push((byte>>i) & 1 == 1);
            }
        }

        let mut bit_idx = 0;
        let mut right_col = self.size - 1;
        let mut going_up = true;

        while right_col >= 1 {
            // skip vertical timing line at col 6
            if right_col == 6 {
                right_col -= 1;
            }

            let left_col = right_col - 1;

            // determine row iteration direction
            let rows: Vec<usize> = if going_up {
                (0..self.size).rev().collect()
            } else {
                (0..self.size).collect()
            };

            for r in rows {
                // check right column cell then right column cell
                for &c in &[right_col, left_col] {
                    if !self.grid[r][c].is_reserved {
                        let is_dark = if bit_idx < bits.len() {
                            bits[bit_idx]
                        } else {
                            false
                        };
                        let color = if is_dark { ModuleColor::Dark } else { ModuleColor::Light };

                        // write bits without marking as reserved
                        self.grid[r][c].color = color;
                        bit_idx += 1;
                    }
                }
            }

            // flip direction for the next 2-column pair
            going_up = !going_up;

            if right_col < 2 { break; }
            right_col -= 2;
        }
    }

    // Applies (or removes applied) mask according to given mask_id
    pub fn toggle_mask(&mut self, mask_id: u8) {
        for r in 0..self.size {
            for c in 0..self.size {
                if !self.grid[r][c].is_reserved {
                    if is_mask_bit_set(mask_id, r, c) {
                        self.grid[r][c].color = match self.grid[r][c].color {
                            ModuleColor::Dark   => ModuleColor::Light,
                            ModuleColor::Light  => ModuleColor::Dark,
                        };
                    }
                }
            }
        }
    }

    // Writes the 15-bit format string into the reserved format modules
    pub fn write_format_string(&mut self, formatstr: u16) {
        let size = self.size;

        for i in 0..15 {
            let bit = ((formatstr >> i) & 1) == 1;
            let color = if bit { ModuleColor::Dark } else { ModuleColor::Light };

            // first copy around top-left finder
            let (row, col) = match i {
                0..=5 => (i, 8),
                6 => (7, 8),
                7 => (8, 8),
                8 => (8, 7),
                9..=14 => (8, 14 - i),
                _ => unreachable!(),
            };
            self.grid[row][col].color = color;

            // second copy split across bottom-left and top-right finders
            if i < 8 {
                // bottom-left (vertical)
                self.grid[8][size-1-i].color = color;
            } else {
                // top-right (horizontal)
                self.grid[size-15+i][8].color = color;
            }
        }
    }

    // Evaluate and select best out of 8 masks
    pub fn select_mask(&mut self, ec_level: &str) {
        let mut best_mask = 0;
        let mut min_penalty = u32::MAX;

        for mask_id in 0..8 {
            // generate and write candidate format string
            let formatstr = generate_format_string(ec_level, mask_id);
            self.write_format_string(formatstr);

            // toggle candidate mask pattern
            self.toggle_mask(mask_id);

            // compute penalty score
            let penalty = self.calculate_penalty();
            if penalty < min_penalty {
                min_penalty = penalty;
                best_mask = mask_id;
            }

            // unapply candidate mask
            self.toggle_mask(mask_id);
        }

        let winning_format = generate_format_string(ec_level, best_mask);
        self.write_format_string(winning_format);
        self.toggle_mask(best_mask);
    }

    // Render ASCII
    pub fn render_ascii(&self) -> String {
        let mut output = String::new();
        let mut border = 4;

        // top quiet zone
        for _ in 0..border {
            for _ in 0..(self.size+border*2) {
                output.push_str("  ");
            }
            output.push_str("\n");
        }

        // QR Grid Rows
        for r in 0..self.size {
            // left quiet zone
            for _ in 0..border {
                output.push_str("  ");
            }

            for c in 0..self.size {
                // If terminal background is dark, Dark modules = "  " (space)
                // and Light modules = "██" (block) for proper contrast.
                match self.grid[r][c].color {
                    ModuleColor::Dark => output.push_str("  "),
                    ModuleColor::Light => output.push_str("██"),
                }
            }

            // right quiet zone
            for _ in 0..border {
                output.push_str("  ");
            }
            output.push('\n');
        }

        // bottom quiet zone
        for _ in 0..border {
            for _ in 0..(self.size + border * 2) {
                output.push_str("  ");
            }
            output.push('\n');
        }

        output
    }

    // Render image (supports png, jpg, etc based on path extension)
    // pub fn render_img(&self, path: &str, scale: u32, border: usize) -> Result<(), image::ImageError> {
    //     let total_modules = self.size + border * 2;
    //     let img_size = (total_modules as u32) * scale;
    //
    //     // gray image buffer filled with white (255) for background/quiet zones
    //     let mut img_buf = ImageBuffer::from_fn(img_size, img_size, |_x, _y| {
    //         Luma([255u8])
    //     });
    //
    //     // fill dark modules
    //     for r in 0..self.size {
    //         for c in 0..self.size {
    //             if self.grid[r][c].color == ModuleColor::Dark {
    //                 let start_x = ((c + border) as u32) * scale;
    //                 let start_y = ((r + border) as u32) * scale;
    //
    //                 for py in start_y..(start_y+scale) {
    //                     for px in start_x..(start_x+scale) {
    //                         img_buf.put_pixel(px, py, Luma([0u8]));
    //                     }
    //                 }
    //             }
    //         }
    //     }
    //
    //     img_buf.save(path)
    // }

    // HELPER: Set specific values at a chosen cell
    pub fn set(&mut self, row: usize, col: usize, color: ModuleColor, is_reserved: bool) {
        if row < self.size && col < self.size {
            self.grid[row][col] = Module { color, is_reserved };
        }
    }

    // HELPER: Calculate penalty across the whole grid
    pub fn calculate_penalty(&self) -> u32 {
        self.penalty_1() + self.penalty_2() + self.penalty_3() + self.penalty_4()
    }

    // PENALTY: Rule 1 - Same color runs of 5+ in a row/col
    fn penalty_1(&self) -> u32 {
        let mut score = 0;

        // horizontal
        for r in 0..self.size {
            let mut run_len = 1;
            for c in 1..self.size {
                if self.grid[r][c].color == self.grid[r][c-1].color {
                    run_len += 1;
                } else {
                    if run_len >= 5 {
                        score += 3 + (run_len - 5);
                    }
                    run_len = 1;
                }
            }
            if run_len >= 5 {
                score += 3 + (run_len - 5);
            }
        }

        // vertical
        for c in 0..self.size {
            let mut run_len = 1;
            for r in 1..self.size {
                if self.grid[r][c].color == self.grid[r-1][c].color {
                    run_len += 1;
                } else {
                    if run_len >= 5 {
                        score += 3 + (run_len - 5);
                    }
                    run_len = 1;
                }
            }
            if run_len >= 5 {
                score += 3 + (run_len - 5);
            }
        }

        score
    }

    // PENALTY: Rule 2 - 2x2 blocks of same color
    fn penalty_2(&self) -> u32 {
        let mut score = 0;

        for r in 0..self.size-1 {
            for c in 0..self.size-1 {
                let color = self.grid[r][c].color;
                if self.grid[r+1][c].color == color
                    && self.grid[r][c+1].color == color
                    && self.grid[r+1][c+1].color == color {
                    score += 3;
                }
            }
        }

        score
    }

    // PENALTY: Rule 3 - False positives for finders
    fn penalty_3(&self) -> u32 {
        let mut score = 0;
        let pattern1 = [
            ModuleColor::Light, ModuleColor::Light, ModuleColor::Light, ModuleColor::Light,
            ModuleColor::Dark, ModuleColor::Light, ModuleColor::Dark, ModuleColor::Dark,
            ModuleColor::Dark, ModuleColor::Light, ModuleColor::Dark,
        ];
        let pattern2 = [
            ModuleColor::Dark, ModuleColor::Light, ModuleColor::Dark, ModuleColor::Dark,
            ModuleColor::Dark, ModuleColor::Light, ModuleColor::Dark, ModuleColor::Light,
            ModuleColor::Light, ModuleColor::Light, ModuleColor::Light,
        ];

        if self.size < 11 {
            return 0;
        }

        // check rows
        for r in 0..self.size {
            for c in 0..=self.size-11 {
                let row_slice: Vec<ModuleColor> = (0..11).map(|i| self.grid[r][c+i].color).collect();
                if row_slice == pattern1 || row_slice == pattern2 {
                    score += 40;
                }
            }
        }

        // check columns
        for c in 0..self.size {
            for r in 0..=self.size-11 {
                let col_slice: Vec<ModuleColor> = (0..11).map(|i| self.grid[r+i][c].color).collect();
                if col_slice == pattern1 || col_slice == pattern2 {
                    score += 40;
                }
            }
        }

        score
    }

    // PENALTY: Rule 4 - Proportion of dark modules
    fn penalty_4(&self) -> u32 {
        let total_modules = (self.size * self.size) as u32;
        let mut dark_count = 0;

        for r in 0..self.size {
            for c in 0..self.size {
                if self.grid[r][c].color == ModuleColor::Dark { dark_count += 1; }
            }
        }

        let percent = (dark_count * 100) / total_modules;
        let diff = if percent >= 50 {
            percent - 50
        } else {
            50 - percent
        };

        let rating = diff / 5;
        rating * 10
    }

    // PATTERN: Finder and separator patterns (7x7 square with 1 row+col of separators)
    fn place_finders_and_separators(&mut self) {
        let size = self.size;

        // Draw finders
        self.draw_finder(0, 0);         // top-left
        self.draw_finder(0, size - 7);    // top-right
        self.draw_finder(size - 7, 0);    // bottom-left

        // Draw separators (Light 1-module border around finders)
        for i in 0..8 {
            // Top-Left Separator (Row 7 and Col 7)
            self.set(7, i, ModuleColor::Light, true);
            self.set(i, 7, ModuleColor::Light, true);

            // Top-Right Separator
            self.set(7, size - 1 - i, ModuleColor::Light, true);
            self.set(i, size - 8, ModuleColor::Light, true);

            // Bottom-Left Separator
            self.set(size - 8, i, ModuleColor::Light, true);
            self.set(size - 1 - i, 7, ModuleColor::Light, true);
        }
    }

    fn draw_finder(&mut self, top_row: usize, left_col: usize) {
        for r in 0..7 {
            for c in 0..7 {
                let is_dark = r == 0 || r == 6 || c == 0 || c == 6 || (r >= 2 && r <= 4 && c >= 2 && c <= 4);
                let color = if is_dark { ModuleColor::Dark } else { ModuleColor::Light };
                self.set(top_row+r, left_col+c, color, true);
            }
        }
    }

    // PATTERN: Timing patterns (alternating light/dark row & col at idx 6)
    fn place_timing_patterns(&mut self) {
        for i in 8..(self.size-8) {
            let color = if i%2 == 0 { ModuleColor::Dark } else { ModuleColor::Light };
            // horizontal
            if !self.grid[6][i].is_reserved {
                self.set(6, i, color, true);
            }
            // vertical
            if !self.grid[i][6].is_reserved {
                self.set(i, 6, color, true);
            }
        }
    }

    // PATTERN: Dark module (always present at row = 4*version + 9, col = 8)
    fn place_dark_module(&mut self) {
        self.set((4*self.version + 9), 8, ModuleColor::Dark, true);
    }

    // PATTERN: Reserve format string modules (around finders)
    fn reserve_format_areas(&mut self) {
        let size = self.size;

        // around top-left finder
        for i in 0..9 {
            if i != 6 {
                self.set(8, i, ModuleColor::Light, true);
                self.set(i, 8, ModuleColor::Light, true);
            }
        }

        // around top-right finder
        for col in (size-8)..size {
            self.set(8, col, ModuleColor::Light, true);
        }

        // around bottom-left finder
        for row in (size-7)..size {
            self.set(row, 8, ModuleColor::Light, true);
        }
    }

    // PATTERN: Alignment patterns (5x5 grid, required for v2+)
    fn place_alignment_patterns(&mut self) {
        if self.version < 2 { return; }

        let coords = get_alignment_coords(self.version);
        for &r in &coords {
            for &c in &coords {
                if self.grid[r][c].is_reserved {
                    continue;
                }
                self.draw_alignment(r, c);
            }
        }
    }

    fn draw_alignment(&mut self, center_r: usize, center_c: usize) {
        for r in -2..=2 {
            for c in -2..=2 {
                let is_dark = (r as isize).abs() == 2 || (c as isize).abs() == 2 || (r==0 && c==0);
                let color = if is_dark { ModuleColor::Dark } else { ModuleColor::Light };
                let row = (center_r as isize + r) as usize;
                let col = (center_c as isize + c) as usize;
                self.set(row, col, color, true);
            }
        }
    }
}



// TESTS
#[cfg(test)]
mod tests {
    use crate::{gf28::GF28, mode_encoder::prepare_data, reed::{interleave_blocks, ReedSolomon}};

    use super::*;

    #[test]
    fn test_v1_matrix_static_patterns() {
        let matrix = Matrix::new(1); // Version 1 -> 21x21
        assert_eq!(matrix.size, 21);

        // Check Top-Left Finder center module is Dark and Reserved
        assert_eq!(matrix.grid[3][3].color, ModuleColor::Dark);
        assert!(matrix.grid[3][3].is_reserved);

        // Check Timing Pattern alternating bits
        assert_eq!(matrix.grid[6][8].color, ModuleColor::Dark);
        assert_eq!(matrix.grid[6][9].color, ModuleColor::Light);
        assert_eq!(matrix.grid[6][10].color, ModuleColor::Dark);

        // Check Dark Module
        assert_eq!(matrix.grid[13][8].color, ModuleColor::Dark);
        assert!(matrix.grid[13][8].is_reserved);

        // Check unreserved data space (e.g., bottom-right corner)
        assert!(!matrix.grid[20][20].is_reserved);
    }

    #[test]
    fn test_data_placement_v1() {
        let mut matrix = Matrix::new(1);
        
        // Sample dummy bytes (e.g., 26 codewords for V1-M)
        let dummy_data = vec![0b10101010; 26];
        matrix.place_data_bits(&dummy_data);

        // Verify bottom-right module (first cell visited in zigzag path: row 20, col 20)
        // Bit 0 of 0b10101010 is 1 (Dark)
        assert_eq!(matrix.grid[20][20].color, ModuleColor::Dark);
        assert!(!matrix.grid[20][20].is_reserved);

        // Verify adjacent module (row 20, col 19)
        // Bit 1 of 0b10101010 is 0 (Light)
        assert_eq!(matrix.grid[20][19].color, ModuleColor::Light);
        assert!(!matrix.grid[20][19].is_reserved);
    }

    // #[test]
    // fn test_render_scannable_qr() {
    //     // use crate::encoder::encode_data;
    //     // use crate::interleave::interleave_blocks;
    //     // use crate::reed::ReedSolomon;
    //     // use crate::gf28::GF28;
    //     // use crate::qr_spec::get_block_spec;
    //
    //     let text = "https://github.com/aether-flux";
    //     let (buf, spec) = prepare_data(text, None).unwrap();
    //
    //     let gf = GF28::new();
    //     let rs = ReedSolomon::new(&gf);
    //     let final_codewords = interleave_blocks(&buf.to_bytes(), &spec, &rs);
    // }
}