use image::RgbaImage;
use ort::session::{Session, builder::GraphOptimizationLevel};
use ort::value::Value;
use std::path::Path;

#[derive(Clone, Debug)]
pub struct DetectedBox {
    pub x: f64,
    pub y: f64,
    pub width: f64,
    pub height: f64,
}

pub fn run_ocr_detection(
    model_path: &Path,
    input_image: &RgbaImage,
) -> Result<Vec<DetectedBox>, String> {
    // 1. Load Model
    let mut session = Session::builder()
        .map_err(|e| format!("Failed to create session: {}", e))?
        .with_optimization_level(GraphOptimizationLevel::Level3)
        .map_err(|e| format!("Failed to set opt: {}", e))?
        .with_intra_threads(4)
        .map_err(|e| format!("Failed to set threads: {}", e))?
        .commit_from_file(model_path)
        .map_err(|e| format!("Failed to load OCR model: {}", e))?;

    // 2. Preprocess
    // DBNet expects standard normalization: (img - mean) / std
    // Mean: [0.485, 0.456, 0.406], Std: [0.229, 0.224, 0.225]
    // And usually resized to multiple of 32. Limit max size for speed.
    let max_side = 1600; // Increase resolution limit
    let (orig_w, orig_h) = input_image.dimensions();
    
    let mut resize_w = orig_w;
    let mut resize_h = orig_h;
    
    // Resize logic: Limit max side, preserve aspect, ensure divisible by 32
    if resize_w > max_side || resize_h > max_side {
        let ratio = max_side as f64 / (orig_w.max(orig_h) as f64);
        resize_w = (orig_w as f64 * ratio) as u32;
        resize_h = (orig_h as f64 * ratio) as u32;
    }
    
    // Align to 32
    resize_w = (resize_w + 31) / 32 * 32;
    resize_h = (resize_h + 31) / 32 * 32;
    
    // Minimum size
    resize_w = resize_w.max(32);
    resize_h = resize_h.max(32);

    let resized = image::imageops::resize(input_image, resize_w, resize_h, image::imageops::FilterType::Triangle);
    
    let channel_stride = (resize_w * resize_h) as usize;
    let mut input_data = Vec::with_capacity(1 * 3 * channel_stride);
    let mut r_plane = Vec::with_capacity(channel_stride);
    let mut g_plane = Vec::with_capacity(channel_stride);
    let mut b_plane = Vec::with_capacity(channel_stride);

    let mean = [0.485, 0.456, 0.406];
    let std = [0.229, 0.224, 0.225];

    for (_x, _y, pixel) in resized.enumerate_pixels() {
        let r = pixel[0] as f32 / 255.0;
        let g = pixel[1] as f32 / 255.0;
        let b = pixel[2] as f32 / 255.0;
        
        r_plane.push((r - mean[0]) / std[0]);
        g_plane.push((g - mean[1]) / std[1]);
        b_plane.push((b - mean[2]) / std[2]);
    }
    input_data.extend(r_plane);
    input_data.extend(g_plane);
    input_data.extend(b_plane);

    // 3. Inference
    let input_shape = vec![1, 3, resize_h as i64, resize_w as i64];
    let input_value = Value::from_array((input_shape, input_data))
        .map_err(|e| format!("Failed to create input tensor: {}", e))?;

    let inputs = ort::inputs![input_value]; // For PP-OCR, usually just one input "x"
    let outputs = session.run(inputs).map_err(|e| format!("Inference failed: {}", e))?;
    
    let output_tensor = outputs.values().next().ok_or("No output")?;
    let (shape, data) = output_tensor.try_extract_tensor::<f32>()
        .map_err(|e| format!("Failed to extract output: {}", e))?;

    // 4. Post-process
    // Output shape is [1, 1, H, W] probability map
    if shape.len() < 4 {
        return Err("Unexpected output shape".to_string());
    }
    
    let map_w = shape[3] as u32;
    let map_h = shape[2] as u32;
    
    // Create binary map (threshold 0.3)
    let threshold = 0.3; 
    let mut binary_map = vec![false; (map_w * map_h) as usize];
    
    for i in 0..binary_map.len() {
        if data[i] > threshold {
            binary_map[i] = true;
        }
    }

    // Find Connected Components (Simple Bounding Box finding)
    let mut visited = vec![false; binary_map.len()];
    let mut boxes = Vec::new();
    
    for y in 0..map_h {
        for x in 0..map_w {
            let idx = (y * map_w + x) as usize;
            if binary_map[idx] && !visited[idx] {
                // Flood fill
                let mut stack = vec![(x, y)];
                visited[idx] = true;
                
                let mut min_x = x;
                let mut max_x = x;
                let mut min_y = y;
                let mut max_y = y;
                let mut pixel_count = 0;
                
                while let Some((cx, cy)) = stack.pop() {
                    pixel_count += 1;
                    if cx < min_x { min_x = cx; }
                    if cx > max_x { max_x = cx; }
                    if cy < min_y { min_y = cy; }
                    if cy > max_y { max_y = cy; }
                    
                    let neighbors = [
                        (cx.wrapping_sub(1), cy), (cx + 1, cy),
                        (cx, cy.wrapping_sub(1)), (cx, cy + 1)
                    ];
                    
                    for (nx, ny) in neighbors {
                        if nx < map_w && ny < map_h {
                            let nidx = (ny * map_w + nx) as usize;
                            if binary_map[nidx] && !visited[nidx] {
                                visited[nidx] = true;
                                stack.push((nx, ny));
                            }
                        }
                    }
                }
                
                // Filter small noise
                if pixel_count < 10 { continue; }
                
                // Calculate Scale Factors
                let scale_x = orig_w as f64 / resize_w as f64;
                let scale_y = orig_h as f64 / resize_h as f64;
                
                // Map to raw coordinates in map space
                let raw_w = (max_x - min_x + 1) as f64;
                let raw_h = (max_y - min_y + 1) as f64;
                
                // --- ASPECT RATIO FILTERING ---
                // Watermarks are typically horizontal text lines. 
                // A cross or vertical pillar will have a small width/height ratio.
                let aspect_ratio = raw_w / raw_h;
                if aspect_ratio < 1.5 {
                    continue; // Skip vertical or square-ish non-text objects
                }

                // --- PADDING / DILATION ---
                let pad_x = raw_w * 0.15; // 15% horizontal is usually enough
                let pad_y = raw_h * 1.00; // Increased to 100% for aggressive vertical coverage
                
                let box_x = (min_x as f64 - pad_x).max(0.0);
                let box_y = (min_y as f64 - pad_y).max(0.0);
                let box_w = raw_w + 2.0 * pad_x;
                let box_h = raw_h + 2.0 * pad_y;

                // Convert to Normalized Image Coordinates [0, 1]
                boxes.push(DetectedBox {
                    x: (box_x * scale_x) / orig_w as f64,
                    y: (box_y * scale_y) / orig_h as f64,
                    width: (box_w * scale_x) / orig_w as f64,
                    height: (box_h * scale_y) / orig_h as f64,
                });
            }
        }
    }

    Ok(boxes)
}
