import os
import subprocess
from pathlib import Path

from garaga.curves import CurveID, ProofSystem
from garaga.modulo_circuit_structs import G2Line, StructArray
from garaga.precompiled_circuits.multi_miller_loop import precompute_lines
from garaga.starknet.cli.utils import create_directory, get_package_version
from garaga.starknet.constants import CAIRO_VERSION, STARKNET_FOUNDRY_VERSION
from garaga.starknet.groth16_contract_generator.calldata import (
    groth16_calldata_from_vk_and_proof,
)
from garaga.starknet.groth16_contract_generator.parsing_utils import (
    Groth16Proof,
    Groth16VerifyingKey,
)

ECIP_OPS_CLASS_HASH = 0x312D1DD5F967EAF6F86965E3FA7ACBC9D0FBD979066A17721DD913736AF2F5E


def precompute_lines_from_vk(vk: Groth16VerifyingKey) -> StructArray:

    # Precompute lines for fixed G2 points
    lines = precompute_lines([vk.gamma, vk.delta])
    precomputed_lines = StructArray(
        name="lines",
        elmts=[
            G2Line(name=f"line{i}", elmts=lines[i : i + 4])
            for i in range(0, len(lines), 4)
        ],
    )

    return precomputed_lines


def gen_test_file(
    contract_name: str,
    verifier_name: str,
    system: ProofSystem,
    verification_function_name: str,
    package_name: str,
    rpc_url: str = "https://rpc.starknet-testnet.lava.build:443",
):
    code = f"""
//! # Starknet Verification Contract Tests
//!
//! This module contains comprehensive tests for the proof verification contract.
//! The tests use Starknet Foundry's testing features to verify proofs.
//!
//! ## Test Strategy
//! - Use real proof fixtures generated by the garaga CLI
//! - Test both positive (valid proof) and negative (invalid proof) cases

use snforge_std::fs::{{FileTrait, read_txt}};
use snforge_std::{{DeclareResultTrait, declare}};
use {package_name}::{verifier_name}::{{
    I{contract_name}LibraryDispatcher, I{contract_name}DispatcherTrait,
}};
use starknet::ClassHash;

/// Declare the {contract_name} contract for testing.
///
/// This function declares a new instance of the {contract_name} contract
///
/// # Arguments
/// * `name` - The contract name to declare
///
/// # Returns
/// * `ClassHash` - The class hash of the declared contract
///
/// # Example
/// ```cairo
/// let class_hash = declare_contract("{contract_name}");
/// let dispatcher = I{contract_name}LibraryDispatcher {{ class_hash }};
/// ```
fn declare_contract(name: ByteArray) -> ClassHash {{
    let class_hash = *declare(name).unwrap().contract_class().class_hash;
    class_hash
}}

/// Test the verification function {verification_function_name} with a valid proof.
///
/// This test demonstrates the complete verification flow:
/// 1. Declare the verification contract
/// 2. Load a pre-generated proof fixture
/// 3. Verify the proof on-chain
/// 4. Assert that verification succeeds
///
/// ## Test Setup
/// - Uses fork testing against Sepolia testnet to access required declared contracts maintained by the garaga team
/// - Loads proof calldata from `tests/proof_calldata.txt`
/// - To generate the calldata, you can use the garaga CLI:
/// ```bash
/// garaga calldata
///     --system {system.value}
///     --vk path/to/vk
///     --proof path/to/proof
///     --public-inputs path/to/public_inputs
///     --format snforge
///     --output-path path/to/tests
/// ```
/// This will generate a file `proof_calldata.txt` in the specified output path.
/// You can then use this file in the test.
///
/// ## Expected Behavior
/// - The proof should verify successfully
/// - The function should return Ok(public_inputs) containing the expected public inputs
/// - No panics or assertion failures should occur
#[test]
#[fork(url: "{rpc_url}", block_tag: latest)]
fn test_{verification_function_name}() {{
    // Step 1: Declare the verification contract
    let class_hash = declare_contract("{contract_name}");
    let dispatcher = I{contract_name}LibraryDispatcher {{ class_hash }};

    // Step 2: Load the proof fixture generated by the script
    // This file contains the {system.value} proof calldata formatted for Starknet
    let file = FileTrait::new("tests/proof_calldata.txt");
    let calldata = read_txt(@file).span();

    // Step 3: Verify the proof on-chain
    // This calls the contract's {verification_function_name} function with the loaded proof data
    let result = dispatcher.{verification_function_name}(calldata);

    // Step 4: Assert that verification succeeded
    // A valid proof should return Ok(public_inputs), not Err
    assert(result.is_ok(), 'Proof verification failed');

    // Optional: You can also verify the public inputs contain expected values
    // let public_inputs = result.unwrap();
    // assert(public_inputs.len() > 0, 'No public inputs');
    // Additional assertions on the public input values can be added here
}}
"""
    return code


def get_scarb_toml_file(package_name: str, cli_mode: bool, inlining_level: int = 2):
    version = get_package_version()
    header = f"# Generated by Garaga version {version}\n"
    if version == "dev":
        suffix = ""
    else:
        suffix = ', tag = "v' + version + '"'
    if cli_mode:
        dep = 'git = "https://github.com/keep-starknet-strange/garaga.git"' + suffix
    else:
        dep = 'path = "../../../"'

    return (
        header
        + f"""[package]
name = "{package_name}"
version = "0.1.0"
edition = "2024_07"

[dependencies]
garaga = {{ {dep} }}
starknet = "{CAIRO_VERSION}"

[cairo]
sierra-replace-ids = false
inlining-strategy = {inlining_level}

[dev-dependencies]
snforge_std = "{STARKNET_FOUNDRY_VERSION}"
assert_macros = "{CAIRO_VERSION}"

[[target.starknet-contract]]
casm = true
casm-add-pythonic-hints = true
"""
    )


def write_test_calldata_file(
    output_folder_path: str,
    vk: Groth16VerifyingKey,
    proof: Groth16Proof,
) -> None:
    """Write the test calldata file for a given verifying key and proof."""
    tests_dir = os.path.join(output_folder_path, "tests")
    cd = groth16_calldata_from_vk_and_proof(vk, proof)
    with open(os.path.join(tests_dir, "proof_calldata.txt"), "w") as f:
        for x in cd[1:]:
            f.write(f"{hex(x)}\n")


def create_verifier_directory_structure(
    output_folder_path: str, include_test_sample: bool = True
) -> tuple[str, str]:
    """Create the standard directory structure for a verifier contract.

    Returns:
        tuple: (src_dir, tests_dir) paths
    """
    create_directory(output_folder_path)
    src_dir = os.path.join(output_folder_path, "src")
    tests_dir = (
        os.path.join(output_folder_path, "tests") if include_test_sample else None
    )
    create_directory(src_dir)
    if include_test_sample:
        create_directory(tests_dir)
    return src_dir, tests_dir


def write_tool_versions_file(output_folder_path: str) -> None:
    """Write the .tool-versions file."""
    with open(os.path.join(output_folder_path, ".tool-versions"), "w") as f:
        f.write(f"scarb {CAIRO_VERSION}\n")
        f.write(f"starknet-foundry {STARKNET_FOUNDRY_VERSION}\n")


def write_lib_cairo_file(src_dir: str, modules: list[str] = None) -> None:
    """Write the lib.cairo file with specified modules."""
    if modules is None:
        modules = ["groth16_verifier", "groth16_verifier_constants"]

    module_imports = "\n".join(f"pub mod {module};" for module in modules)
    with open(os.path.join(src_dir, "lib.cairo"), "w") as f:
        version = get_package_version()
        header = f"// Generated by Garaga version {version}\n"
        f.write(header + f"\n{module_imports}\n")


def write_verifier_files(
    output_folder_path: str,
    package_name: str,
    constants_code: str | None,
    contract_code: str,
    contract_cairo_name: str,
    verification_function_name: str,
    system: ProofSystem,
    cli_mode: bool = False,
    circuits_code: str | None = None,
    modules: list[str] = None,
    constants_filename: str = "groth16_verifier_constants.cairo",
    contract_filename: str = "groth16_verifier.cairo",
    circuits_filename: str | None = None,
    include_test_sample: bool = True,
    constants_output_path: str | None = None,
) -> None:
    """Write all the standard verifier files (directories, constants, contract, scarb.toml, lib.cairo, test files)."""

    # Create directory structure
    src_dir, tests_dir = create_verifier_directory_structure(
        output_folder_path, include_test_sample
    )

    # Write tool versions
    write_tool_versions_file(output_folder_path)

    # Build header once for Cairo files
    version = get_package_version()
    header_cairo = f"// Generated by Garaga version {version}\n"

    # Write Scarb.toml
    with open(os.path.join(output_folder_path, "Scarb.toml"), "w") as f:
        f.write(get_scarb_toml_file(package_name, cli_mode))

    # Write lib.cairo
    write_lib_cairo_file(src_dir, modules)

    # Write constants file (if requested)
    if constants_code is not None:
        target_constants_path = constants_output_path or os.path.join(
            src_dir, constants_filename
        )
        target_constants_dir = os.path.dirname(target_constants_path)
        os.makedirs(target_constants_dir, exist_ok=True)
        with open(target_constants_path, "w") as f:
            f.write(header_cairo + constants_code)

        subprocess.run(["scarb", "fmt", "."], cwd=target_constants_dir, check=True)

    # Write contract file
    with open(os.path.join(src_dir, contract_filename), "w") as f:
        f.write(header_cairo + contract_code)

    # Write circuits file if provided
    if circuits_code and circuits_filename:
        with open(os.path.join(src_dir, circuits_filename), "w") as f:
            f.write(header_cairo + circuits_code)

    # Write test file.
    if include_test_sample:
        with open(os.path.join(tests_dir, "test_contract.cairo"), "w") as f:
            f.write(
                header_cairo
                + gen_test_file(
                    contract_cairo_name,
                    contract_filename.removesuffix(".cairo"),
                    system,
                    verification_function_name,
                    package_name,
                )
            )
    subprocess.run(["scarb", "fmt", "."], cwd=output_folder_path, check=True)


def write_test_calldata_file_generic(
    output_folder_path: str,
    system: ProofSystem,
    vk_path: str,
    proof_path: str,
    public_inputs_path: str = None,
) -> None:
    """Write the test calldata file for any proof system using get_calldata_generic."""
    from pathlib import Path as PathType

    from garaga.starknet.cli.verify import get_calldata_generic

    tests_dir = os.path.join(output_folder_path, "tests")
    cd = get_calldata_generic(
        system,
        PathType(vk_path),
        PathType(proof_path),
        PathType(public_inputs_path) if public_inputs_path else None,
    )
    with open(os.path.join(tests_dir, "proof_calldata.txt"), "w") as f:
        for x in cd[1:]:
            f.write(f"{hex(x)}\n")


def gen_groth16_verifier(
    vk: str | Path | Groth16VerifyingKey,
    output_folder_path: str,
    output_folder_name: str,
    ecip_class_hash: int = ECIP_OPS_CLASS_HASH,
    cli_mode: bool = False,
    include_test_sample: bool = True,
) -> str:
    if isinstance(vk, (Path, str)):
        vk = Groth16VerifyingKey.from_json(vk)
    else:
        vk = vk

    curve_id = vk.curve_id
    if cli_mode:
        output_folder_name = output_folder_name
    else:
        output_folder_name = output_folder_name + f"_{curve_id.name.lower()}"
    output_folder_path = os.path.join(output_folder_path, output_folder_name)

    precomputed_lines = precompute_lines_from_vk(vk)
    verification_function_name = f"verify_groth16_proof_{curve_id.name.lower()}"
    contract_cairo_name = f"Groth16Verifier{curve_id.name}"
    constants_code = f"""
    use garaga::definitions::{{G1Point, G2Point, E12D, G2Line, u384}};
    {f"use garaga::definitions::u288;" if curve_id!=CurveID.BLS12_381 else ""}
    use garaga::groth16::Groth16VerifyingKey;

    pub const N_PUBLIC_INPUTS:usize = {len(vk.ic)-1};
    {vk.serialize_to_cairo()}
    pub const precomputed_lines: [G2Line; {len(precomputed_lines)//4}] = {precomputed_lines.serialize(raw=True, const=True)};
    """
    contract_code = f"""
use super::groth16_verifier_constants::{{N_PUBLIC_INPUTS, vk, ic, precomputed_lines}};

#[starknet::interface]
pub trait I{contract_cairo_name}<TContractState> {{
    fn {verification_function_name}(
        self: @TContractState,
        full_proof_with_hints: Span<felt252>,
    ) -> Result<Span<u256>, felt252>;
}}

#[starknet::contract]
mod {contract_cairo_name} {{
    use starknet::SyscallResultTrait;
    use garaga::definitions::{{G1Point, G1G2Pair}};
    use garaga::groth16::{{multi_pairing_check_{curve_id.name.lower()}_3P_2F_with_extra_miller_loop_result, Groth16ProofRawTrait}};
    use garaga::ec_ops::{{G1PointTrait, ec_safe_add}};
    use garaga::utils::calldata::{{deserialize_full_proof_with_hints_{curve_id.name.lower()}}};
    use super::{{N_PUBLIC_INPUTS, vk, ic, precomputed_lines}};

    const ECIP_OPS_CLASS_HASH: felt252 = {hex(ecip_class_hash)};

    #[storage]
    struct Storage {{}}

    #[abi(embed_v0)]
    impl IGroth16Verifier{curve_id.name} of super::IGroth16Verifier{curve_id.name}<ContractState> {{
        fn {verification_function_name}(
            self: @ContractState,
            full_proof_with_hints: Span<felt252>,
        ) -> Result<Span<u256>, felt252> {{
            // DO NOT EDIT THIS FUNCTION UNLESS YOU KNOW WHAT YOU ARE DOING.
            // This function returns Result::Ok(public_inputs) if the proof is valid.
            // If the proof is invalid, it returns Result::Err(error).
            // Read the documentation to learn how to generate the full_proof_with_hints array given a proof and a verifying key.
            let fph = deserialize_full_proof_with_hints_{curve_id.name.lower()}(full_proof_with_hints);
            let groth16_proof = fph.groth16_proof;
            let mpcheck_hint = fph.mpcheck_hint;
            let msm_hint = fph.msm_hint;

            groth16_proof.raw.check_proof_points({curve_id.value});

            let ic = ic.span();

            let vk_x: G1Point = match ic.len() {{
                0 => panic!("Malformed VK"),
                1 => *ic.at(0),
                _ => {{
                    // Start serialization with the hint array directly to avoid copying it.
                    let mut msm_calldata: Array<felt252> = array![];
                    // Add the points from VK and public inputs to the proof.
                    Serde::serialize(@ic.slice(1, N_PUBLIC_INPUTS), ref msm_calldata);
                    Serde::serialize(@groth16_proof.public_inputs, ref msm_calldata);
                    // Complete with the curve indentifier ({curve_id.value} for {curve_id.name}):
                    msm_calldata.append({curve_id.value});
                    // Add the hint array.
                    for x in msm_hint {{
                        msm_calldata.append(*x);
                    }}

                    // Call the multi scalar multiplication endpoint on the Garaga ECIP ops contract
                    // to obtain vk_x.
                    let mut _vx_x_serialized = starknet::syscalls::library_call_syscall(
                        ECIP_OPS_CLASS_HASH.try_into().unwrap(),
                        selector!("msm_g1"),
                        msm_calldata.span()
                    )
                        .unwrap_syscall();

                    ec_safe_add(
                        Serde::<G1Point>::deserialize(ref _vx_x_serialized).unwrap(), *ic.at(0), {curve_id.value}
                    )
                }}
            }};
            // Perform the pairing check.
            let check = multi_pairing_check_{curve_id.name.lower()}_3P_2F_with_extra_miller_loop_result(
                G1G2Pair {{ p: vk_x, q: vk.gamma_g2 }},
                G1G2Pair {{ p: groth16_proof.raw.c, q: vk.delta_g2 }},
                G1G2Pair {{ p: groth16_proof.raw.a.negate({curve_id.value}), q: groth16_proof.raw.b }},
                vk.alpha_beta_miller_loop_result,
                precomputed_lines.span(),
                mpcheck_hint,
            );
            match check {{
                Result::Ok(_) => Result::Ok(groth16_proof.public_inputs),
                Result::Err(error) => Result::Err(error),
            }}
        }}
    }}
}}


    """

    # Use the new reusable function to write all files
    write_verifier_files(
        output_folder_path,
        output_folder_name,
        constants_code,
        contract_code,
        contract_cairo_name,
        verification_function_name,
        ProofSystem.Groth16,
        cli_mode,
        include_test_sample=include_test_sample,
    )

    return constants_code


if __name__ == "__main__":

    EXAMPLES_FOLDER = "hydra/garaga/starknet/groth16_contract_generator/examples"
    vk_paths = [
        os.path.join(EXAMPLES_FOLDER, "vk_bn254.json"),
        os.path.join(EXAMPLES_FOLDER, "vk_bls.json"),
    ]
    proof_paths = [
        os.path.join(EXAMPLES_FOLDER, "proof_bn254.json"),
        os.path.join(EXAMPLES_FOLDER, "proof_bls.json"),
    ]

    CONTRACTS_FOLDER = "src/contracts/autogenerated/"  # Do not change this

    FOLDER_NAME = "groth16_example"  # '_curve_id' is appended in the end.

    def _generate_verifier(vk_path):
        try:
            gen_groth16_verifier(
                vk_path, CONTRACTS_FOLDER, FOLDER_NAME, ECIP_OPS_CLASS_HASH
            )
        except Exception as e:
            print(f"An error occurred: {e}")
            raise e

    for vk_path in vk_paths:
        _generate_verifier(vk_path)

    FOLDER_NAMES = [f"groth16_example_bn254", "groth16_example_bls12_381"]
    for contract_name, vk_path, proof_path in zip(FOLDER_NAMES, vk_paths, proof_paths):
        vk = Groth16VerifyingKey.from_json(vk_path)
        proof = Groth16Proof.from_json(proof_path=proof_path)
        output_folder_path = os.path.join(CONTRACTS_FOLDER, contract_name)
        write_test_calldata_file_generic(
            output_folder_path,
            system=ProofSystem.Groth16,
            vk_path=vk_path,
            proof_path=proof_path,
        )
