use std::fs;
use std::io::{Read, Seek, SeekFrom};

use clap::{Parser, Subcommand};
use pyaxml_rs::arsc::Arsc;
use pyaxml_rs::axml::Axml;

/// Android binary XML (AXML) and resource table (ARSC) converter.
#[derive(Parser)]
#[command(name = "pyaxml", version, about, long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Convert binary AXML to plain XML.
    Axml2xml {
        /// Input file (binary AXML or APK/ZIP archive)
        #[arg(short, long)]
        input: String,

        /// Output file (defaults to stdout)
        #[arg(short, long)]
        output: Option<String>,

        /// Entry path inside APK/ZIP (default: AndroidManifest.xml)
        #[arg(short, long)]
        path: Option<String>,

        /// Export string pool to JSON file, or import from it
        #[arg(long)]
        stringblocks_file: Option<String>,
    },

    /// Convert plain XML to binary AXML.
    Xml2axml {
        /// Input XML file
        #[arg(short, long)]
        input: String,

        /// Output binary AXML file
        #[arg(short, long)]
        output: String,

        /// Import string pool from JSON file
        #[arg(long)]
        stringblocks_file: Option<String>,
    },

    /// Convert ARSC resource table to XML, grouped by locale.
    Arsc2xml {
        /// Input file (binary ARSC or APK/ZIP archive)
        #[arg(short, long)]
        input: String,

        /// Output file (defaults to stdout)
        #[arg(short, long)]
        output: Option<String>,

        /// Entry path inside APK/ZIP (default: resources.arsc)
        #[arg(short, long)]
        path: Option<String>,

        /// Filter by locale tag: 'default' for no-locale entries, 'en', 'fr-FR', etc.
        #[arg(short, long)]
        language: Option<String>,
    },

    /// Convert ARSC resource table to protobuf text format.
    Arsc2proto {
        /// Input file (binary ARSC or APK/ZIP archive)
        #[arg(short, long)]
        input: String,

        /// Output file (defaults to stdout)
        #[arg(short, long)]
        output: Option<String>,

        /// Entry path inside APK/ZIP (default: resources.arsc)
        #[arg(short, long)]
        path: Option<String>,

        /// Pretty-print output with one field per line and indentation
        #[arg(long)]
        pretty: bool,
    },

    /// Convert binary AXML to protobuf text format.
    Axml2proto {
        /// Input file (binary AXML or APK/ZIP archive)
        #[arg(short, long)]
        input: String,

        /// Output file (defaults to stdout)
        #[arg(short, long)]
        output: Option<String>,

        /// Entry path inside APK/ZIP (default: AndroidManifest.xml)
        #[arg(short, long)]
        path: Option<String>,

        /// Pretty-print output with one field per line and indentation
        #[arg(long)]
        pretty: bool,
    },

    /// Parse a binary AXML or ARSC file and write it as a binary proto file.
    ExportProto {
        /// Input file (binary AXML or binary ARSC)
        #[arg(short, long)]
        input: String,

        /// Output binary proto file
        #[arg(short, long)]
        output: String,

        /// Force type: 'axml' or 'arsc' (auto-detected from magic bytes when omitted)
        #[arg(long)]
        r#type: Option<String>,
    },

    /// Load a binary proto file and emit XML text or a re-packed binary.
    ImportProto {
        /// Input binary proto file (produced by export-proto)
        #[arg(short, long)]
        input: String,

        /// Write XML text to this file (or stdout if omitted)
        #[arg(long)]
        to_xml: Option<String>,

        /// Write re-packed binary to this file
        #[arg(long)]
        pack: Option<String>,

        /// Force type: 'axml' or 'arsc' (auto-detected when omitted; tries AXML first)
        #[arg(long)]
        r#type: Option<String>,
    },
}

/// Maximum uncompressed size accepted from a ZIP entry (256 MiB).
/// Prevents ZIP-bomb decompression attacks (SEC-04).
const MAX_ZIP_ENTRY_BYTES: u64 = 256 * 1024 * 1024;

/// Read binary data from a plain file or a named entry inside a ZIP/APK.
///
/// The file is opened exactly once regardless of whether it turns out to be a ZIP (P7).
/// `default_entry` is used as the ZIP entry name when `path` is `None`.
fn read_binary_data(
    input: &str,
    path: Option<&str>,
    default_entry: &str,
) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
    let mut file = fs::File::open(input)?;

    // Peek at the first two bytes to detect ZIP without buffering the whole file.
    let mut magic = [0u8; 2];
    let n = file.read(&mut magic)?;

    if n == 2 && &magic == b"PK" {
        // Seek back so ZipArchive can read the full archive from the start.
        file.seek(SeekFrom::Start(0))?;
        let target = path.unwrap_or(default_entry);

        // SEC-05: reject paths that escape the archive root.
        if target.contains("..") || target.starts_with('/') || target.starts_with('\\') {
            return Err(format!("Invalid ZIP entry path: '{}'", target).into());
        }

        let mut archive = zip::ZipArchive::new(file)?;
        let inner = match archive.by_name(target) {
            Ok(f) => f,
            Err(_) => return Err(format!("'{}' not found inside '{}'", target, input).into()),
        };

        // SEC-04: stream at most MAX_ZIP_ENTRY_BYTES+1 so we detect bombs even
        // when the ZIP central-directory uncompressed_size field is forged.
        let mut buf = Vec::new();
        inner.take(MAX_ZIP_ENTRY_BYTES + 1).read_to_end(&mut buf)?;
        if buf.len() as u64 > MAX_ZIP_ENTRY_BYTES {
            return Err(format!(
                "ZIP entry '{}' exceeds maximum decompression size ({} MiB)",
                target,
                MAX_ZIP_ENTRY_BYTES / 1024 / 1024
            )
            .into());
        }
        Ok(buf)
    } else {
        let mut buf = magic[..n].to_vec();
        file.read_to_end(&mut buf)?;
        Ok(buf)
    }
}

/// Read binary AXML data from a plain file or from a named entry inside a ZIP/APK.
fn read_axml_data(input: &str, path: Option<&str>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
    read_binary_data(input, path, "AndroidManifest.xml")
}

/// Read binary ARSC data from a plain file or from `resources.arsc` inside a ZIP/APK.
fn read_arsc_data(input: &str, path: Option<&str>) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
    read_binary_data(input, path, "resources.arsc")
}

/// Write text to a file or stdout.
fn write_output(output: Option<&str>, data: &str) -> Result<(), Box<dyn std::error::Error>> {
    match output {
        Some(path) => fs::write(path, data)?,
        None => print!("{}", data),
    }
    Ok(())
}

/// Write the string pool as JSON to `path`.
fn write_stringblocks_file(axml: &Axml, path: &str) -> Result<(), Box<dyn std::error::Error>> {
    fs::write(path, axml.export_stringblocks_json())?;
    Ok(())
}

/// Pre-populate the string pool from a JSON file.
fn read_stringblocks_file(path: &str, axml: &mut Axml) -> Result<(), Box<dyn std::error::Error>> {
    let content = fs::read_to_string(path)?;
    axml.import_stringblocks_json(&content);
    Ok(())
}

fn die(msg: impl std::fmt::Display) -> ! {
    eprintln!("pyaxml: {}", msg);
    std::process::exit(1);
}

fn main() {
    let cli = Cli::parse();

    match cli.command {
        Command::Axml2xml {
            input,
            output,
            path,
            stringblocks_file,
        } => {
            let data = read_axml_data(&input, path.as_deref()).unwrap_or_else(|e| die(e));
            let axml = Axml::from_axml(&data).unwrap_or_else(|e| die(e));
            let xml = axml.to_xml().unwrap_or_else(|e| die(e));
            write_output(output.as_deref(), &xml).unwrap_or_else(|e| die(e));
            if let Some(ref sb) = stringblocks_file {
                write_stringblocks_file(&axml, sb).unwrap_or_else(|e| {
                    eprintln!("pyaxml: warning: could not write stringblocks file: {}", e);
                });
            }
        }

        Command::Axml2proto {
            input,
            output,
            path,
            pretty,
        } => {
            let data = read_axml_data(&input, path.as_deref()).unwrap_or_else(|e| die(e));
            let axml = Axml::from_axml(&data).unwrap_or_else(|e| die(e));
            let proto_text = if pretty {
                axml.to_proto_text_pretty()
            } else {
                axml.to_proto_text()
            };
            write_output(output.as_deref(), &proto_text).unwrap_or_else(|e| die(e));
        }

        Command::Xml2axml {
            input,
            output,
            stringblocks_file,
        } => {
            let xml = fs::read_to_string(&input).unwrap_or_else(|e| die(e));
            let mut axml = Axml::new();
            if let Some(ref sb) = stringblocks_file {
                read_stringblocks_file(sb, &mut axml).unwrap_or_else(|e| {
                    eprintln!("pyaxml: warning: could not read stringblocks file: {}", e);
                });
            }
            axml.from_xml(&xml)
                .unwrap_or_else(|e: pyaxml_rs::AxmlError| die(e));
            fs::write(&output, axml.pack()).unwrap_or_else(|e| die(e));
        }

        Command::Arsc2proto {
            input,
            output,
            path,
            pretty,
        } => {
            let data = read_arsc_data(&input, path.as_deref()).unwrap_or_else(|e| die(e));
            let arsc = Arsc::from_axml(&data).unwrap_or_else(|e| die(e));
            let proto_text = if pretty {
                arsc.to_proto_text_pretty()
            } else {
                arsc.to_proto_text()
            };
            write_output(output.as_deref(), &proto_text).unwrap_or_else(|e| die(e));
        }

        Command::Arsc2xml {
            input,
            output,
            path,
            language,
        } => {
            let data = read_arsc_data(&input, path.as_deref()).unwrap_or_else(|e| die(e));
            let arsc = Arsc::from_axml(&data).unwrap_or_else(|e| die(e));
            let xml = arsc.list_packages(language.as_deref());
            write_output(output.as_deref(), &xml).unwrap_or_else(|e| die(e));
        }

        Command::ExportProto {
            input,
            output,
            r#type,
        } => {
            let is_arsc = match r#type.as_deref() {
                Some("arsc") => true,
                Some("axml") => false,
                _ => {
                    // Detect from magic bytes: ARSC files start with 0x02 0x00.
                    let magic = {
                        let mut f = fs::File::open(&input).unwrap_or_else(|e| die(e));
                        let mut b = [0u8; 2];
                        f.read_exact(&mut b).unwrap_or(());
                        b
                    };
                    magic == [0x02, 0x00]
                }
            };

            if is_arsc {
                let data = read_arsc_data(&input, None).unwrap_or_else(|e| die(e));
                let arsc = Arsc::from_axml(&data).unwrap_or_else(|e| die(e));
                arsc.export_proto_file(&output).unwrap_or_else(|e| die(e));
            } else {
                let data = read_axml_data(&input, None).unwrap_or_else(|e| die(e));
                let axml = Axml::from_axml(&data).unwrap_or_else(|e| die(e));
                axml.export_proto_file(&output).unwrap_or_else(|e| die(e));
            }
        }

        Command::ImportProto {
            input,
            to_xml,
            pack,
            r#type,
        } => {
            let raw = fs::read(&input).unwrap_or_else(|e| die(e));

            // Detect type: try AXML decode first unless --type is given.
            let use_arsc = match r#type.as_deref() {
                Some("arsc") => true,
                Some("axml") => false,
                _ => {
                    // Try AXML; fall back to ARSC on failure.
                    Axml::from_proto_bytes(&raw).is_err()
                }
            };

            if use_arsc {
                let arsc = Arsc::from_proto_bytes(&raw).unwrap_or_else(|e| die(e));
                if let Some(ref path) = pack {
                    fs::write(path, arsc.pack()).unwrap_or_else(|e| die(e));
                }
                if to_xml.is_some() || pack.is_none() {
                    let xml = arsc.list_packages(None);
                    write_output(to_xml.as_deref(), &xml).unwrap_or_else(|e| die(e));
                }
            } else {
                let axml = Axml::from_proto_bytes(&raw).unwrap_or_else(|e| die(e));
                if let Some(ref path) = pack {
                    fs::write(path, axml.pack()).unwrap_or_else(|e| die(e));
                }
                if to_xml.is_some() || pack.is_none() {
                    let xml = axml.to_xml().unwrap_or_else(|e| die(e));
                    write_output(to_xml.as_deref(), &xml).unwrap_or_else(|e| die(e));
                }
            }
        }
    }
}
