//! Form field value write-back and flattening
//!
//! Provides FormFieldWriter for updating form field values in raw PDF bytes,
//! flattening forms, and exporting form data as FDF/XFDF.

use crate::fitz::error::{Error, Result};
use crate::fitz::geometry::Rect;
use crate::pdf::form::{Form, FormField, WidgetType};
use crate::pdf::write::{escape_pdf_name, escape_pdf_string};
use std::collections::HashMap;

/// Writer for form field values into raw PDF bytes.
///
/// Takes a mutable reference to raw PDF bytes and form fields.
/// Updates /V and /AS entries in-place when encoded length matches.
pub struct FormFieldWriter<'a> {
    pdf_bytes: &'a mut [u8],
    /// Map from field name to (obj_num, dict_start, dict_end)
    field_locations: HashMap<String, (i32, usize, usize)>,
}

impl<'a> FormFieldWriter<'a> {
    /// Create a writer from mutable PDF bytes and a form.
    /// Scans the PDF to build field name -> object location map.
    pub fn new(pdf_bytes: &'a mut [u8], form: &Form) -> Result<Self> {
        let field_locations = scan_field_locations(pdf_bytes)?;
        Ok(Self {
            pdf_bytes,
            field_locations,
        })
    }

    /// Update /V entry in the field's PDF dictionary.
    /// Returns error if the field is not found or if the new value's encoded
    /// length differs from the existing value (in-place replacement only).
    pub fn set_field_value_in_pdf(&mut self, field: &FormField, value: &str) -> Result<()> {
        let (obj_num, dict_start, dict_end) = self
            .field_locations
            .get(field.name())
            .copied()
            .ok_or_else(|| Error::argument(format!("Field '{}' not found in PDF", field.name())))?;

        let region = &self.pdf_bytes[dict_start..dict_end];
        let v_pos = find_dict_key(region, b"/V")
            .ok_or_else(|| Error::argument(format!("Field '{}' has no /V entry", field.name())))?;
        let abs_v_pos = dict_start + v_pos;

        let (value_start, value_end, is_name) = find_value_extent(self.pdf_bytes, abs_v_pos)?;
        let old_len = value_end - value_start;

        let new_encoded = if is_name {
            format!("/{}", escape_pdf_name(value))
        } else {
            let escaped = escape_pdf_string(value.as_bytes());
            format!("({})", String::from_utf8_lossy(&escaped))
        };
        let new_bytes = new_encoded.as_bytes();

        if new_bytes.len() != old_len {
            return Err(Error::argument(format!(
                "New value encoded length {} differs from existing {} (in-place replacement only)",
                new_bytes.len(),
                old_len
            )));
        }

        self.pdf_bytes[value_start..value_end].copy_from_slice(new_bytes);
        Ok(())
    }

    /// Update /AS entry for checkbox/radio (appearance state).
    /// Sets /AS to /Yes when checked, /Off when unchecked.
    pub fn set_field_checked_in_pdf(&mut self, field: &FormField, checked: bool) -> Result<()> {
        if !matches!(
            field.field_type(),
            WidgetType::Checkbox | WidgetType::RadioButton
        ) {
            return Err(Error::argument(
                "set_field_checked_in_pdf only applies to checkbox/radio fields",
            ));
        }

        let (_, dict_start, dict_end) = self
            .field_locations
            .get(field.name())
            .copied()
            .ok_or_else(|| Error::argument(format!("Field '{}' not found in PDF", field.name())))?;

        let region = &self.pdf_bytes[dict_start..dict_end];
        let as_pos = find_dict_key(region, b"/AS").ok_or_else(|| {
            Error::argument(format!(
                "Field '{}' has no /AS entry (widget may need /AP)",
                field.name()
            ))
        })?;
        let abs_as_pos = dict_start + as_pos;

        let (value_start, value_end, _) = find_value_extent(self.pdf_bytes, abs_as_pos)?;
        let old_len = value_end - value_start;

        let new_value = if checked { "/Yes" } else { "/Off" };
        if new_value.len() != old_len {
            return Err(Error::argument(format!(
                "/AS value length mismatch: need {} bytes, have {}",
                new_value.len(),
                old_len
            )));
        }

        self.pdf_bytes[value_start..value_end].copy_from_slice(new_value.as_bytes());
        Ok(())
    }

    /// Generate minimal appearance stream (/AP /N) for text fields with the current value.
    /// Creates a simple BT...Tj...ET stream.
    pub fn update_appearance_stream(&mut self, _field: &FormField) -> Result<()> {
        // Minimal implementation: generate BT 1 0 0 1 x y Tm (value) Tj ET
        // For full implementation we'd need to locate/create the /AP/N stream
        // and replace its content. This requires more PDF structure parsing.
        Err(Error::unsupported(
            "update_appearance_stream requires full PDF object parsing",
        ))
    }

    /// Flatten form: render field values into page content, remove widget annotations,
    /// remove /AcroForm from Catalog, remove field objects.
    pub fn flatten_form(&mut self, _form: &Form) -> Result<()> {
        // Full flattening requires:
        // 1. For each page with widget annotations, get page Contents stream(s)
        // 2. Prepend content that draws each field's value (BT...Tj...ET)
        // 3. Remove widget refs from page /Annots
        // 4. Remove /AcroForm from Catalog
        // 5. Remove field objects (or leave as orphaned - cleaner to remove)
        Err(Error::unsupported(
            "flatten_form requires page/contents/annot parsing",
        ))
    }

    /// Export form data as FDF format.
    pub fn export_fdf(form: &Form, source_filename: Option<&str>) -> Result<Vec<u8>> {
        let mut out = Vec::new();
        out.extend_from_slice(b"%FDF-1.2\n");
        out.extend_from_slice(&[0x25, 0xE2, 0xE3, 0xCF, 0xD3, 0x0A]); // PDF binary comment
        out.extend_from_slice(b"1 0 obj\n<< /FDF << /Fields [\n");

        for field in form.fields() {
            let name_escaped = escape_pdf_string(field.name().as_bytes());
            let value_escaped = escape_pdf_string(field.value().as_bytes());
            let entry = format!(
                "<< /T ({}) /V ({}) >>\n",
                String::from_utf8_lossy(&name_escaped),
                String::from_utf8_lossy(&value_escaped)
            );
            out.extend_from_slice(entry.as_bytes());
        }

        out.extend_from_slice(b"] >>");
        if let Some(f) = source_filename {
            let f_escaped = escape_pdf_string(f.as_bytes());
            out.extend_from_slice(
                format!(" /F ({}) ", String::from_utf8_lossy(&f_escaped)).as_bytes(),
            );
        }
        out.extend_from_slice(b">> >>\nendobj\ntrailer\n<< /Root 1 0 R >>\n%%EOF\n");
        Ok(out)
    }

    /// Export form data as XFDF (XML) format.
    pub fn export_xfdf(form: &Form) -> Result<Vec<u8>> {
        let mut out = Vec::new();
        out.extend_from_slice(b"<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n");
        out.extend_from_slice(
            b"<xfdf xmlns=\"http://ns.adobe.com/xfdf/\" xml:space=\"preserve\">\n",
        );
        out.extend_from_slice(b"  <fields>\n");

        for field in form.fields() {
            let name_escaped = escape_xml(field.name());
            let value_escaped = escape_xml(field.value());
            out.extend_from_slice(
                format!(
                    "    <field name=\"{}\">\n      <value>{}</value>\n    </field>\n",
                    name_escaped, value_escaped
                )
                .as_bytes(),
            );
        }

        out.extend_from_slice(b"  </fields>\n</xfdf>\n");
        Ok(out)
    }
}

fn escape_xml(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for c in s.chars() {
        match c {
            '&' => out.push_str("&amp;"),
            '<' => out.push_str("&lt;"),
            '>' => out.push_str("&gt;"),
            '"' => out.push_str("&quot;"),
            '\'' => out.push_str("&apos;"),
            _ => out.push(c),
        }
    }
    out
}

/// Scan PDF for /AcroForm /Fields and build field name -> (obj_num, dict_start, dict_end).
fn scan_field_locations(data: &[u8]) -> Result<HashMap<String, (i32, usize, usize)>> {
    let text = String::from_utf8_lossy(data);
    let mut map = HashMap::new();

    let acroform_start = text
        .find("/AcroForm")
        .ok_or_else(|| Error::format("No /AcroForm in PDF"))?;
    let search_region = &text[acroform_start..];
    let dict_start_rel = search_region
        .find("<<")
        .ok_or_else(|| Error::format("No dict after /AcroForm"))?;
    let dict_start = acroform_start + dict_start_rel;

    let dict_text = extract_balanced_dict(&text, dict_start)
        .ok_or_else(|| Error::format("Unbalanced dict at AcroForm"))?;

    let fields_pos = dict_text
        .find("/Fields")
        .ok_or_else(|| Error::format("No /Fields in AcroForm"))?;
    let after_fields = &dict_text[fields_pos + 7..];
    let arr_start = after_fields
        .find('[')
        .ok_or_else(|| Error::format("No array after /Fields"))?;
    let arr_end = after_fields[arr_start..]
        .find(']')
        .ok_or_else(|| Error::format("Unclosed /Fields array"))?;
    let arr_text = &after_fields[arr_start + 1..arr_start + arr_end];

    let obj_refs: Vec<&str> = arr_text
        .split('R')
        .filter(|s| !s.trim().is_empty())
        .collect();

    for obj_ref in obj_refs {
        let parts: Vec<&str> = obj_ref.split_whitespace().collect();
        if parts.len() >= 2 {
            if let Ok(obj_num) = parts[0].parse::<i32>() {
                if let Some((name, dict_start_abs, dict_end_abs)) =
                    find_field_dict_and_name(data, obj_num)
                {
                    map.insert(name, (obj_num, dict_start_abs, dict_end_abs));
                }
            }
        }
    }

    Ok(map)
}

fn extract_balanced_dict(text: &str, start: usize) -> Option<String> {
    let bytes = text.as_bytes();
    if start + 1 >= bytes.len() || bytes[start] != b'<' || bytes[start + 1] != b'<' {
        return None;
    }
    let mut depth = 0i32;
    let mut i = start;
    while i < bytes.len().saturating_sub(1) {
        if bytes[i] == b'<' && bytes[i + 1] == b'<' {
            depth += 1;
            i += 2;
        } else if bytes[i] == b'>' && bytes[i + 1] == b'>' {
            depth -= 1;
            i += 2;
            if depth == 0 {
                return Some(text[start..i].to_string());
            }
        } else {
            i += 1;
        }
    }
    Some(text[start..].to_string())
}

fn find_field_dict_and_name(data: &[u8], obj_num: i32) -> Option<(String, usize, usize)> {
    let pattern = format!("{} 0 obj", obj_num);
    let pat_bytes = pattern.as_bytes();
    let pos = data.windows(pat_bytes.len()).position(|w| w == pat_bytes)?;
    let after_obj = &data[pos + pat_bytes.len()..];
    let dict_rel = after_obj.windows(2).position(|w| w == b"<<")?;
    let dict_start = pos + pat_bytes.len() + dict_rel;

    let dict_end = find_dict_end(data, dict_start)?;
    let end_inclusive = (dict_end + 2).min(data.len());
    let dict_region = std::str::from_utf8(&data[dict_start..end_inclusive]).ok()?;

    let name = extract_pdf_string_from_dict(dict_region, "/T")?;
    Some((name, dict_start, dict_end + 2))
}

fn find_dict_end(data: &[u8], start: usize) -> Option<usize> {
    let mut depth = 0i32;
    let mut i = start;
    while i < data.len().saturating_sub(1) {
        if data[i] == b'<' && data[i + 1] == b'<' {
            depth += 1;
            i += 2;
        } else if data[i] == b'>' && data[i + 1] == b'>' {
            depth -= 1;
            i += 2;
            if depth == 0 {
                return Some(i - 2);
            }
        } else {
            i += 1;
        }
    }
    None
}

fn find_dict_key(region: &[u8], key: &[u8]) -> Option<usize> {
    region
        .windows(key.len())
        .position(|w| w == key)
        .map(|p| p + key.len())
}

fn extract_pdf_string_from_dict(dict: &str, key: &str) -> Option<String> {
    // Find key as whole token: /T not /FT. Require /T followed by space or (
    let search = format!("{} ", key);
    let pos = dict.find(&search).or_else(|| {
        let search2 = format!("{}(", key);
        dict.find(&search2).map(|p| p)
    })?;
    let after = dict[pos + key.len()..].trim_start();
    if after.starts_with('(') {
        let mut depth = 0i32;
        let mut end = 0;
        for (i, ch) in after.chars().enumerate() {
            match ch {
                '(' => depth += 1,
                ')' => {
                    depth -= 1;
                    if depth == 0 {
                        end = i;
                        break;
                    }
                }
                _ => {}
            }
        }
        if end > 1 {
            Some(after[1..end].to_string())
        } else {
            None
        }
    } else if after.starts_with('/') {
        let name_end = after[1..]
            .find(|c: char| {
                c.is_whitespace() || c == '/' || c == '>' || c == '<' || c == '[' || c == '('
            })
            .map(|i| i + 1)
            .unwrap_or(after.len());
        Some(after[1..name_end].to_string())
    } else {
        None
    }
}

/// Find the extent of the value after a dict key (position after the key).
/// Returns (value_start, value_end, is_name).
fn find_value_extent(data: &[u8], pos_after_key: usize) -> Result<(usize, usize, bool)> {
    let mut i = pos_after_key;
    while i < data.len() && data[i].is_ascii_whitespace() {
        i += 1;
    }
    if i >= data.len() {
        return Err(Error::format("Unexpected end after dict key"));
    }
    let value_start = i;
    if data[i] == b'(' {
        let mut depth = 0i32;
        while i < data.len() {
            match data[i] {
                b'(' => {
                    depth += 1;
                    i += 1;
                }
                b')' => {
                    depth -= 1;
                    i += 1;
                    if depth == 0 {
                        return Ok((value_start, i, false));
                    }
                }
                b'\\' if i + 1 < data.len() => {
                    i += 2;
                }
                _ => i += 1,
            }
        }
        Err(Error::format("Unclosed string in PDF"))
    } else if data[i] == b'/' {
        i += 1; // skip leading / of name
        while i < data.len()
            && !data[i].is_ascii_whitespace()
            && data[i] != b'/'
            && data[i] != b'>'
            && data[i] != b']'
        {
            i += 1;
        }
        Ok((value_start, i, true))
    } else if data[i] == b'<' && i + 1 < data.len() && data[i + 1] != b'<' {
        i += 1;
        while i < data.len() && data[i] != b'>' {
            if data[i] == b'#' && i + 2 < data.len() {
                i += 3;
            } else {
                i += 1;
            }
        }
        if i < data.len() {
            i += 1;
        }
        Ok((value_start, i, false))
    } else {
        Err(Error::format("Unknown value type after dict key"))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_simple_pdf_with_form() -> Vec<u8> {
        let s = r#"%PDF-1.4
1 0 obj
<< /Type /Catalog /Pages 2 0 R /AcroForm << /Fields [5 0 R] >> >>
endobj
2 0 obj
<< /Type /Pages /Count 1 /Kids [3 0 R] >>
endobj
3 0 obj
<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Annots [5 0 R] >>
endobj
5 0 obj
<< /Type /Annot /Subtype /Widget /FT /Tx /T (username) /V (oldvalue) /Rect [100 700 300 730] >>
endobj
xref
0 6
0000000000 65535 f
0000000009 00000 n
0000000058 00000 n
0000000115 00000 n
0000000172 00000 n
0000000229 00000 n
trailer << /Size 6 /Root 1 0 R >>
startxref
300
%%EOF"#;
        s.as_bytes().to_vec()
    }

    fn make_pdf_with_checkbox() -> Vec<u8> {
        let s = r#"%PDF-1.4
1 0 obj
<< /Type /Catalog /Pages 2 0 R /AcroForm << /Fields [5 0 R] >> >>
endobj
2 0 obj
<< /Type /Pages /Count 1 /Kids [3 0 R] >>
endobj
3 0 obj
<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>
endobj
5 0 obj
<< /Type /Annot /Subtype /Widget /FT /Btn /T (agree) /V /Off /AS /Off /Rect [100 700 120 720] >>
endobj
xref
0 6
0000000000 65535 f
0000000009 00000 n
0000000058 00000 n
0000000115 00000 n
0000000172 00000 n
0000000229 00000 n
trailer << /Size 6 /Root 1 0 R >>
startxref
300
%%EOF"#;
        s.as_bytes().to_vec()
    }

    #[test]
    fn test_export_fdf() {
        let mut form = Form::new();
        let mut f1 = FormField::text_field("name".to_string(), Rect::EMPTY, None);
        f1.set_value("John".to_string()).unwrap();
        form.add_field(f1);
        let mut f2 = FormField::text_field("email".to_string(), Rect::EMPTY, None);
        f2.set_value("j@x.com".to_string()).unwrap();
        form.add_field(f2);

        let fdf = FormFieldWriter::export_fdf(&form, None).unwrap();
        let fdf_str = String::from_utf8_lossy(&fdf);
        assert!(fdf_str.contains("%FDF-1.2"));
        assert!(fdf_str.contains("/Fields"));
        assert!(fdf_str.contains("name"));
        assert!(fdf_str.contains("John"));
        assert!(fdf_str.contains("email"));
        assert!(fdf_str.contains("j@x.com"));
        assert!(fdf_str.contains("%%EOF"));
    }

    #[test]
    fn test_export_fdf_with_filename() {
        let mut form = Form::new();
        let mut f = FormField::text_field("x".to_string(), Rect::EMPTY, None);
        f.set_value("y".to_string()).unwrap();
        form.add_field(f);

        let fdf = FormFieldWriter::export_fdf(&form, Some("form.pdf")).unwrap();
        let fdf_str = String::from_utf8_lossy(&fdf);
        assert!(fdf_str.contains("form.pdf"));
    }

    #[test]
    fn test_export_xfdf() {
        let mut form = Form::new();
        let mut f1 = FormField::text_field("field1".to_string(), Rect::EMPTY, None);
        f1.set_value("value1".to_string()).unwrap();
        form.add_field(f1);

        let xfdf = FormFieldWriter::export_xfdf(&form).unwrap();
        let xfdf_str = String::from_utf8_lossy(&xfdf);
        assert!(xfdf_str.contains("<?xml"));
        assert!(xfdf_str.contains("xfdf"));
        assert!(xfdf_str.contains("field1"));
        assert!(xfdf_str.contains("value1"));
    }

    #[test]
    fn test_export_xfdf_escapes_xml() {
        let mut form = Form::new();
        let mut f = FormField::text_field("x".to_string(), Rect::EMPTY, None);
        f.set_value("a<b>c&d".to_string()).unwrap();
        form.add_field(f);

        let xfdf = FormFieldWriter::export_xfdf(&form).unwrap();
        let xfdf_str = String::from_utf8_lossy(&xfdf);
        assert!(xfdf_str.contains("&lt;"));
        assert!(xfdf_str.contains("&gt;"));
        assert!(xfdf_str.contains("&amp;"));
    }

    #[test]
    fn test_scan_field_locations() {
        let pdf = make_simple_pdf_with_form();
        let map = scan_field_locations(&pdf).unwrap();
        assert!(
            map.contains_key("username"),
            "expected 'username' in {:?}",
            map
        );
    }

    #[test]
    fn test_set_field_value_in_pdf_same_length() {
        let mut pdf = make_simple_pdf_with_form();
        let mut form = Form::new();
        let mut field = FormField::text_field("username".to_string(), Rect::EMPTY, None);
        field.set_value("newvalue".to_string()).unwrap();
        form.add_field(field);

        let mut writer = FormFieldWriter::new(&mut pdf, &form).unwrap();
        writer
            .set_field_value_in_pdf(form.get_field("username").unwrap(), "newvalue")
            .unwrap();

        let pdf_str = String::from_utf8_lossy(&pdf);
        assert!(pdf_str.contains("/V (newvalue)"));
    }

    #[test]
    fn test_set_field_value_in_pdf_different_length_fails() {
        let mut pdf = make_simple_pdf_with_form();
        let mut form = Form::new();
        let mut field = FormField::text_field("username".to_string(), Rect::EMPTY, None);
        field.set_value("x".to_string()).unwrap();
        form.add_field(field);

        let mut writer = FormFieldWriter::new(&mut pdf, &form).unwrap();
        let r = writer.set_field_value_in_pdf(form.get_field("username").unwrap(), "longer");
        assert!(r.is_err());
    }

    #[test]
    fn test_set_field_checked_in_pdf() {
        let mut pdf = make_pdf_with_checkbox();
        let mut form = Form::new();
        let mut field = FormField::checkbox("agree".to_string(), Rect::EMPTY, false);
        field.set_checked(true).unwrap();
        form.add_field(field);

        let mut writer = FormFieldWriter::new(&mut pdf, &form).unwrap();
        writer
            .set_field_checked_in_pdf(form.get_field("agree").unwrap(), true)
            .unwrap();

        let pdf_str = String::from_utf8_lossy(&pdf);
        assert!(pdf_str.contains("/AS /Yes"));
    }

    #[test]
    fn test_flatten_form_unsupported() {
        let mut pdf = make_simple_pdf_with_form();
        let mut form = Form::new();
        let field = FormField::text_field("username".to_string(), Rect::EMPTY, None);
        form.add_field(field);

        let mut writer = FormFieldWriter::new(&mut pdf, &form).unwrap();
        let r = writer.flatten_form(&form);
        assert!(r.is_err());
    }

    #[test]
    fn test_escape_xml() {
        assert_eq!(escape_xml("a<b"), "a&lt;b");
        assert_eq!(escape_xml("a>b"), "a&gt;b");
        assert_eq!(escape_xml("a&b"), "a&amp;b");
        assert_eq!(escape_xml("a\"b"), "a&quot;b");
    }
}
