use std::fs;
use std::path::{Path, PathBuf};
use crate::types::Result;
use crate::utils::log::{success, warning, info};
use printpdf::*;
use lopdf::{Document, Object};
use std::io::Cursor;
use std::collections::BTreeMap;
use super::commands::{BundleArgs, PaperKind, Orientation, SortKind};

struct InputItem {
    path: PathBuf,
    rotate: Option<i64>,
}

fn paper_size_mm(paper: &PaperKind, orientation: &Orientation) -> (Mm, Mm) {
    let (w, h) = match paper {
        PaperKind::A4 => (210.0, 297.0),
        PaperKind::Letter => (215.9, 279.4),
        PaperKind::Legal => (215.9, 355.6),
    };
    match orientation {
        Orientation::Portrait => (Mm(w), Mm(h)),
        Orientation::Landscape => (Mm(h), Mm(w)),
    }
}

fn is_supported_image(path: &Path) -> bool {
    if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
        let e = ext.to_lowercase();
        matches!(
            e.as_str(),
            "png" | "jpg" | "jpeg" | "webp" | "gif" | "tiff" | "tif" | "bmp" | "ico" | "qoi" | "pnm" | "tga" | "hdr" | "exr" | "ff" | "avif"
        )
    } else {
        false
    }
}

fn is_pdf(path: &Path) -> bool {
    if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
        ext.to_lowercase() == "pdf"
    } else {
        false
    }
}

fn parse_input(s: &str) -> InputItem {
    // Check for :angle suffix
    // Be careful with Windows paths like C:\...
    // We look for :number at the end.
    if let Some(idx) = s.rfind(':') {
        let suffix = &s[idx+1..];
        if let Ok(angle) = suffix.parse::<i64>() {
            // Check if angle is a common rotation value or just a number
            // Actually, any number is valid for rotation (though usually 90 multiples)
            // But we should ensure it's not part of a path.
            // On unix, : is valid in filename but rare.
            // On windows, : is drive separator (index 1).
            if idx > 1 {
                return InputItem {
                    path: PathBuf::from(&s[..idx]),
                    rotate: Some(angle),
                };
            }
        }
    }
    InputItem {
        path: PathBuf::from(s),
        rotate: None,
    }
}

fn collect_inputs(args: &BundleArgs) -> Vec<InputItem> {
    let mut items: Vec<InputItem> = Vec::new();
    
    // Process explicit inputs
    for input_str in &args.inputs {
        let item = parse_input(input_str);
        if item.path.exists() {
            items.push(item);
        } else {
            warning(&format!("file not found: {}", item.path.display()));
        }
    }

    // Process directory
    if let Some(dir) = &args.dir {
        if dir.is_dir() {
            if let Ok(read) = fs::read_dir(dir) {
                let mut dir_files: Vec<PathBuf> = Vec::new();
                for entry in read.flatten() {
                    let p = entry.path();
                    if p.is_file() && (is_supported_image(&p) || is_pdf(&p)) {
                        dir_files.push(p);
                    }
                }
                
                // Sort dir files
                match args.sort {
                    SortKind::None => {}
                    SortKind::Name => dir_files.sort_by(|a, b| a.file_name().cmp(&b.file_name())),
                }
                
                for p in dir_files {
                    items.push(InputItem {
                        path: p,
                        rotate: None,
                    });
                }
            }
        }
    }
    
    items
}

fn load_dynamic_image(path: &Path) -> Option<printpdf::image_crate::DynamicImage> {
    match printpdf::image_crate::open(path) {
        Ok(img) => Some(img),
        Err(_) => None,
    }
}

fn image_to_pdf(path: &Path, args: &BundleArgs) -> Result<Document> {
    let Some(img) = load_dynamic_image(path) else {
        return Err(format!("failed to load image: {}", path.display()).into());
    };

    let (page_w, page_h) = paper_size_mm(&args.paper, &args.orientation);
    let (doc, page, layer) = PdfDocument::new("image", page_w, page_h, "Layer 1");
    
    let pdf_img = Image::from_dynamic_image(&img);

    let img_w_px = img.width() as f32;
    let img_h_px = img.height() as f32;

    let margin: f32 = 10.0;
    let avail_w: f32 = page_w.0 - margin * 2.0;
    let avail_h: f32 = page_h.0 - margin * 2.0;
    let px_to_mm: f32 = 0.3527;
    let img_w_mm: f32 = img_w_px * px_to_mm;
    let img_h_mm: f32 = img_h_px * px_to_mm;
    let scale: f32 = (avail_w / img_w_mm).min(avail_h / img_h_mm);
    let draw_w: f32 = img_w_mm * scale;
    let draw_h: f32 = img_h_mm * scale;

    let x: f32 = (page_w.0 - draw_w) / 2.0;
    let y: f32 = (page_h.0 - draw_h) / 2.0;

    let layer_ref = doc.get_page(page).get_layer(layer);
    pdf_img.add_to_layer(
        layer_ref,
        ImageTransform {
            translate_x: Some(Mm(x)),
            translate_y: Some(Mm(y)),
            rotate: None,
            scale_x: Some(draw_w / img_w_mm),
            scale_y: Some(draw_h / img_h_mm),
            dpi: None,
        },
    );

    let mut buf = Cursor::new(Vec::new());
    let mut writer = std::io::BufWriter::new(&mut buf);
    doc.save(&mut writer)?;
    drop(writer); // flush
    
    let bytes = buf.into_inner();
    let doc = Document::load_mem(&bytes).map_err(|e| format!("failed to parse generated pdf: {}", e))?;
    Ok(doc)
}

fn apply_rotation(doc: &mut Document, angle: i64) {
    // angle must be multiple of 90?
    // PDF spec says Rotate is integer, usually multiple of 90.
    // We normalize it to 0, 90, 180, 270 (positive).
    let angle = ((angle % 360) + 360) % 360;
    
    for page_id in doc.get_pages().values() {
        if let Ok(page) = doc.get_object_mut(*page_id).and_then(|obj| obj.as_dict_mut()) {
            // If existing rotation exists, add to it? Or replace?
            // "configuration rotation" usually means set to this.
            // But if scanned is sideways, we might want relative?
            // Let's assume absolute setting for now.
            page.set("Rotate", angle);
        }
    }
}

// Helper to renumber object IDs
fn renumber_objects(doc: &mut Document, shift: u32) {
    let mut new_objects = BTreeMap::new();
    for (id, mut object) in std::mem::take(&mut doc.objects) {
        renumber_object(&mut object, shift);
        new_objects.insert((id.0 + shift, id.1), object);
    }
    doc.objects = new_objects;
    
    // Also need to update the trailer / catalog references if they were inside objects?
    // doc.trailer and doc.catalog are dictionaries/references.
    // We need to renumber them too if we were to use them, but we are merging pages.
}

fn renumber_object(obj: &mut Object, shift: u32) {
    match obj {
        Object::Reference(id) => {
            id.0 += shift;
        }
        Object::Array(arr) => {
            for item in arr {
                renumber_object(item, shift);
            }
        }
        Object::Dictionary(dict) => {
            for (_, value) in dict.iter_mut() {
                renumber_object(value, shift);
            }
        }
        Object::Stream(stream) => {
            for (_, value) in stream.dict.iter_mut() {
                renumber_object(value, shift);
            }
        }
        _ => {}
    }
}

fn merge_documents(mut docs: Vec<Document>) -> Result<Document> {
    if docs.is_empty() {
        return Err("no documents to merge".into());
    }
    
    let mut target = docs.remove(0);
    // Ensure target has a Pages root
    // lopdf Document usually has a Catalog pointing to Pages.
    // We will append pages from other docs to target's page tree.
    
    // Simplification: We assume we can append to the root Pages object.
    // Flattening page tree might be needed if it's complex, but for simple merge:
    
    let catalog_id = target.trailer.get(b"Root")
        .and_then(|obj| obj.as_reference())
        .map_err(|_| "Target PDF has no Root")?;
        
    let catalog = target.get_object(catalog_id)
        .and_then(|obj| obj.as_dict())
        .map_err(|_| "Target Root is not a dictionary")?;
        
    let pages_id = catalog.get(b"Pages")
        .and_then(|obj| obj.as_reference())
        .map_err(|_| "Target Catalog has no Pages")?;
    
    // Get current max ID
    let mut max_id = target.max_id;
    
    for (_, mut source) in docs.into_iter().enumerate() {
        // 1. Get pages from source (before renumbering/moving)
        // We need the original IDs to shift them correctly later
        let source_pages = source.get_pages();
        
        // 2. Renumber source
        let shift = max_id + 1;
        renumber_objects(&mut source, shift);
        max_id += source.max_id + 1; // Update max_id
        
        // 3. Add source objects to target
        for (id, object) in source.objects {
            target.objects.insert(id, object);
        }
        
        // Update Parent of the moved pages to point to target's Pages object
        for (_, page_id) in &source_pages {
             let new_id = (page_id.0 + shift, page_id.1);
             if let Ok(page) = target.get_object_mut(new_id).and_then(|o| o.as_dict_mut()) {
                 page.set("Parent", Object::Reference(pages_id));
             }
        }
        
        // 4. Append source pages to target
        // We need to add these page IDs to target's Pages kids
        if let Ok(pages_dict) = target.get_object_mut(pages_id).and_then(|o| o.as_dict_mut()) {
            if let Ok(kids) = pages_dict.get_mut(b"Kids").and_then(|o| o.as_array_mut()) {
                 for (_, page_id) in source_pages {
                     // Shift the page ID manually since we got them before renumbering
                     let new_id = (page_id.0 + shift, page_id.1);
                     kids.push(Object::Reference(new_id));
                 }
            }
        }
    }
    
    target.max_id = max_id;

    // Fix Count in root Pages
    // Note: get_pages() might fail if the structure is inconsistent during modification?
    // Let's rely on manual count.
    // We started with 1 page (usually) or N pages.
    // We added source_pages.len() for each source.
    
    // Better: Count how many kids we have in the root Pages object.
    let mut total_count = 0;
    if let Ok(pages_dict) = target.get_object(pages_id).and_then(|o| o.as_dict()) {
        if let Ok(kids) = pages_dict.get(b"Kids").and_then(|o| o.as_array()) {
            total_count = kids.len() as i64;
        }
    }
    
    // Let's try to trust get_pages() but debug if it fails.
    let pages = target.get_pages();
    let calculated_count = pages.len() as i64;
    
    // If get_pages returns 0, something is wrong.
    if calculated_count == 0 && total_count > 0 {
        // Fallback to kids length if flat
        // But we should warn
        if let Ok(pages_dict) = target.get_object_mut(pages_id).and_then(|o| o.as_dict_mut()) {
            pages_dict.set("Count", total_count);
        }
    } else {
        if let Ok(pages_dict) = target.get_object_mut(pages_id).and_then(|o| o.as_dict_mut()) {
            pages_dict.set("Count", calculated_count);
        }
    }
    
    // Remove structure tree to avoid errors? (Optional)
    // target.trailer.remove(b"Root"); // No, Catalog is Root. Do NOT remove Root!
    // Maybe remove logical structure to keep it simple
    
    Ok(target)
}

pub fn execute(args: BundleArgs) -> Result<()> {
    let items = collect_inputs(&args);
    if items.is_empty() {
        return Err("no input files found".into());
    }
    
    info(&format!("Found {} input files", items.len()));
    
    let mut documents = Vec::new();
    
    for (idx, item) in items.iter().enumerate() {
        info(&format!("[{}/{}] Processing {}", idx+1, items.len(), item.path.display()));
        
        let mut doc = if is_pdf(&item.path) {
            Document::load(&item.path).map_err(|e| format!("failed to load pdf {}: {}", item.path.display(), e))?
        } else if is_supported_image(&item.path) {
            image_to_pdf(&item.path, &args)?
        } else {
            warning(&format!("Skipping unsupported file: {}", item.path.display()));
            continue;
        };
        
        // Apply rotation
        if let Some(angle) = item.rotate.or(args.rotate) {
            apply_rotation(&mut doc, angle);
        }
        
        documents.push(doc);
    }
    
    if documents.is_empty() {
        return Err("No valid documents to bundle".into());
    }
    
    info("Merging documents...");
    let mut merged = merge_documents(documents)?;
    
    // Compress?
    // merged.compress();
    
    info("Saving...");
    merged.save(&args.output).map_err(|e| format!("failed to save pdf: {}", e))?;
    
    success(&format!("saved: {}", args.output.display()));
    Ok(())
}
