mod context;
mod render;

use crate::bytes::move_hex_literal;
use crate::curves::CurveAdapter;
use crate::error::{Error, Result};
use crate::model::Groth16VerifierInputs;
pub use context::{MovegenMode, MovegenTemplateInput};
use handlebars::Handlebars;
use std::env;
use std::fs::{self, create_dir_all, write};
use std::path::{Component, Path};

#[derive(Debug, Clone)]
pub struct GenerateMovePackageOptions<'a> {
    pub package_name: &'a str,
    pub module_name: &'a str,
    pub account_address: &'a str,
    pub mode: MovegenMode,
    pub force: bool,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProofDataSnippet {
    pub proof_a: String,
    pub proof_b: String,
    pub proof_c: String,
    pub public_inputs_rendered: String,
}

impl ProofDataSnippet {
    pub fn render_aptos_test_functions(&self) -> String {
        format!(
            r#"    fun proof_a_bytes(): vector<u8> {{ {} }}
    fun proof_b_bytes(): vector<u8> {{ {} }}
    fun proof_c_bytes(): vector<u8> {{ {} }}
    fun public_inputs_bytes(): vector<vector<u8>> {{ {} }}"#,
            self.proof_a, self.proof_b, self.proof_c, self.public_inputs_rendered
        )
    }
}

pub fn proof_data_snippet(
    adapter: &dyn CurveAdapter,
    inputs: &Groth16VerifierInputs,
) -> Result<ProofDataSnippet> {
    let proof = inputs.proof.as_ref().ok_or_else(|| {
        Error::MissingInput("proof-data requires proof input; VK-only inputs have no proof".into())
    })?;

    let public_inputs_bytes: Vec<String> = inputs
        .public_inputs
        .iter()
        .map(|value| {
            adapter
                .serialize_fr_public_input(value)
                .map(|bytes| move_hex_literal(&bytes))
        })
        .collect::<Result<_>>()?;

    Ok(ProofDataSnippet {
        proof_a: move_hex_literal(&adapter.serialize_g1_proof(&proof.pi_a)?),
        proof_b: move_hex_literal(&adapter.serialize_g2_proof(&proof.pi_b)?),
        proof_c: move_hex_literal(&adapter.serialize_g1_proof(&proof.pi_c)?),
        public_inputs_rendered: render::vector_of_hex(&public_inputs_bytes),
    })
}

pub fn generate_move_package(
    out_dir: &Path,
    adapter: &dyn CurveAdapter,
    inputs: &Groth16VerifierInputs,
    options: &GenerateMovePackageOptions<'_>,
) -> Result<()> {
    validate_account_address(options.account_address)?;
    if options.force {
        validate_safe_force_output_dir(out_dir)?;
    }

    if out_dir.exists() && !options.force {
        return Err(Error::OutputExists(out_dir.to_path_buf()));
    }

    if out_dir.exists() {
        fs::remove_dir_all(out_dir).map_err(|e| Error::Io {
            source: e,
            context: format!("failed to clear existing output dir {}", out_dir.display()),
        })?;
    }

    create_dir_all(out_dir).map_err(|e| Error::Io {
        source: e,
        context: format!("create output dir {}", out_dir.display()),
    })?;

    create_dir_all(out_dir.join("sources")).map_err(|e| Error::Io {
        source: e,
        context: format!("create sources dir {}", out_dir.join("sources").display()),
    })?;

    let mut reg = Handlebars::new();
    register_templates(&mut reg)?;

    let vk = &inputs.verifying_key;
    let public_inputs = &inputs.public_inputs;

    let vk_alpha_g1 = move_hex_literal(&adapter.serialize_g1_vk(&vk.vk_alpha_1)?);
    let vk_beta_g2 = move_hex_literal(&adapter.serialize_g2_vk(&vk.vk_beta_2)?);
    let vk_gamma_g2 = move_hex_literal(&adapter.serialize_g2_vk(&vk.vk_gamma_2)?);
    let vk_delta_g2 = move_hex_literal(&adapter.serialize_g2_vk(&vk.vk_delta_2)?);
    let vk_gamma_abc_g1: Vec<String> = vk
        .ic
        .iter()
        .map(|point| {
            let bytes = adapter.serialize_g1_vk(point)?;
            Ok(move_hex_literal(&bytes))
        })
        .collect::<Result<_>>()?;
    let vk_gamma_abc_g1_rendered = render::vector_of_hex(&vk_gamma_abc_g1);
    let (proof_a, proof_b, proof_c) = match inputs.proof.as_ref() {
        Some(proof) => (
            move_hex_literal(&adapter.serialize_g1_proof(&proof.pi_a)?),
            move_hex_literal(&adapter.serialize_g2_proof(&proof.pi_b)?),
            move_hex_literal(&adapter.serialize_g1_proof(&proof.pi_c)?),
        ),
        None => (
            "x\"\"".to_string(),
            "x\"\"".to_string(),
            "x\"\"".to_string(),
        ),
    };
    let public_inputs_bytes: Vec<String> = public_inputs
        .iter()
        .map(|value| {
            adapter
                .serialize_fr_public_input(value)
                .map(|bytes| move_hex_literal(&bytes))
        })
        .collect::<Result<_>>()?;
    let public_inputs_rendered = render::vector_of_hex(&public_inputs_bytes);
    let invalid_public_inputs = invalid_public_inputs(public_inputs);
    let invalid_public_inputs_bytes: Vec<String> = invalid_public_inputs
        .iter()
        .map(|value| {
            adapter
                .serialize_fr_public_input(value)
                .map(|bytes| move_hex_literal(&bytes))
        })
        .collect::<Result<_>>()?;
    let invalid_public_inputs_rendered = render::vector_of_hex(&invalid_public_inputs_bytes);

    let input = MovegenTemplateInput {
        package_name: options.package_name.to_string(),
        module_name: options.module_name.to_string(),
        account_address: options.account_address.to_string(),
        named_address: options.package_name.to_string(),
        vk_alpha_g1,
        vk_beta_g2,
        vk_gamma_g2,
        vk_delta_g2,
        vk_gamma_abc_g1,
        vk_gamma_abc_g1_rendered,
        proof_a,
        proof_b,
        proof_c,
        public_inputs_bytes,
        public_inputs_rendered,
        invalid_public_inputs_rendered,
        include_entry: options.mode.include_entry(),
    };

    let move_toml = reg
        .render("move_toml", &input)
        .map_err(|e| Error::TemplateRender(e.to_string()))?;
    fs::write(out_dir.join("Move.toml"), move_toml).map_err(|e| Error::Io {
        source: e,
        context: "write Move.toml".to_string(),
    })?;

    let verifier_template = adapter.move_template_name();
    let verifier_source = reg
        .render(verifier_template, &input)
        .map_err(|e| Error::TemplateRender(e.to_string()))?;
    fs::write(
        out_dir.join("sources").join("verifier.move"),
        verifier_source,
    )
    .map_err(|e| Error::Io {
        source: e,
        context: "write verifier.move".to_string(),
    })?;

    if inputs.has_test_vectors() {
        create_dir_all(out_dir.join("tests")).map_err(|e| Error::Io {
            source: e,
            context: format!("create tests dir {}", out_dir.join("tests").display()),
        })?;
        let tests = reg
            .render("move_tests", &input)
            .map_err(|e| Error::TemplateRender(e.to_string()))?;
        fs::write(out_dir.join("tests").join("verifier_tests.move"), tests).map_err(|e| {
            Error::Io {
                source: e,
                context: "write verifier_tests.move".to_string(),
            }
        })?;
    }

    write(
        out_dir.join("README.md"),
        render::readme_content(
            options.package_name,
            options.module_name,
            &input.account_address,
        ),
    )
    .map_err(|e| Error::Io {
        source: e,
        context: "write README.md".to_string(),
    })?;

    Ok(())
}

fn invalid_public_inputs(public_inputs: &[String]) -> Vec<String> {
    let mut invalid = public_inputs.to_vec();
    if let Some(last) = invalid.last_mut() {
        *last = if last == "0" {
            "1".to_string()
        } else {
            "0".to_string()
        };
    }
    invalid
}

fn validate_account_address(value: &str) -> Result<()> {
    let Some(hex) = value
        .strip_prefix("0x")
        .or_else(|| value.strip_prefix("0X"))
    else {
        return Err(Error::InvalidAccountAddress(
            "account_address must start with 0x".to_string(),
        ));
    };
    if hex.is_empty() || hex.len() > 64 || !hex.chars().all(|c| c.is_ascii_hexdigit()) {
        return Err(Error::InvalidAccountAddress(
            "account_address must match 0x[0-9a-fA-F]{1,64}".to_string(),
        ));
    }
    Ok(())
}

fn validate_safe_force_output_dir(out_dir: &Path) -> Result<()> {
    if out_dir
        .components()
        .any(|component| matches!(component, Component::ParentDir))
    {
        return Err(Error::UnsafeOutputDirectory(out_dir.to_path_buf()));
    }

    if !out_dir.exists() {
        return Ok(());
    }

    let target = out_dir.canonicalize().map_err(|e| Error::Io {
        source: e,
        context: format!("canonicalize output dir {}", out_dir.display()),
    })?;
    if target.parent().is_none() {
        return Err(Error::UnsafeOutputDirectory(target));
    }

    let cwd = env::current_dir().map_err(|e| Error::Io {
        source: e,
        context: "get current working directory".to_string(),
    })?;
    let cwd = cwd.canonicalize().map_err(|e| Error::Io {
        source: e,
        context: format!("canonicalize current working directory {}", cwd.display()),
    })?;
    if target == cwd || cwd.starts_with(&target) {
        return Err(Error::UnsafeOutputDirectory(target));
    }

    Ok(())
}

fn register_templates(handlebars: &mut Handlebars) -> Result<()> {
    let move_toml = include_str!("../../templates/Move.toml.hbs");
    let verifier_bn254 = include_str!("../../templates/verifier_bn254.move.hbs");
    let verifier_bls = include_str!("../../templates/verifier_bls12381.move.hbs");
    let tests = include_str!("../../templates/tests.move.hbs");

    handlebars
        .register_template_string("move_toml", move_toml)
        .map_err(|e| Error::TemplateRender(e.to_string()))?;
    handlebars
        .register_template_string("verifier_bn254.move.hbs", verifier_bn254)
        .map_err(|e| Error::TemplateRender(e.to_string()))?;
    handlebars
        .register_template_string("verifier_bls12381.move.hbs", verifier_bls)
        .map_err(|e| Error::TemplateRender(e.to_string()))?;
    handlebars
        .register_template_string("move_tests", tests)
        .map_err(|e| Error::TemplateRender(e.to_string()))?;
    Ok(())
}
