"""
Library containing utility functions related to
EZKL and the EZKL Proof Service.
"""

import json
import logging
import os
import tempfile
from functools import lru_cache
from pathlib import Path
from typing import Any, Dict, Optional

import numpy
import numpy as np
from ezkl import ezkl  # type: ignore
from pydantic import BaseModel

from infernet_ml.resource.artifact_manager import RitualArtifactManager
from infernet_ml.resource.repo_id import RitualRepoId
from infernet_ml.utils.codec.ezkl_codec import extract_io_from_proof_execution
from infernet_ml.utils.codec.vector import RitualVector
from infernet_ml.utils.model_analyzer import ONNXModelAnalyzer  # type: ignore
from infernet_ml.utils.onnx_utils import get_onnx_flops, get_onnx_input_names
from infernet_ml.zk.ezkl.ezkl_artifact import EZKLArtifact
from infernet_ml.zk.ezkl.types import ONNXInput
from infernet_ml.zk.proof_utils import DEFAULT_PROOFS_DIRECTORY, get_proof_dir

log = logging.getLogger(__name__)


class ProofGenerationArgs(BaseModel):
    """
    Arguments for generating the proof.

    Attributes:
        compiled_path (Path): Path to the compiled model
        settings_path (Path): Path to the settings
    """

    compiled_path: Path
    settings_path: Path


class EZKLProof(BaseModel):
    """
    All of the artifacts needed to verify a proof.

    Attributes:
        ezkl_proof (Dict[str, Any]): Proof dictionary as generated by ezkl
        witness (Dict[str, Any]): Witness used to generate the proof
        verify_key (Optional[bytes]): Verification key used to verify the proof.
            Available if a new setup has been done, otherwise None

        input (Optional[List[float]]): Input data
        output (Optional[List[float]]): Output data
        flops (Optional[float]): Number of FLOPS in the model

        ezkl_proof_path (Optional[Path]): Path to the proof file
        witness_path (Optional[Path]): Path to the witness file
        settings_path (Optional[Path]): Path to the settings file
        vk_path (Optional[Path]): Path to the verification
    """

    ezkl_proof: Dict[str, Any]
    witness: Dict[str, Any]

    verify_key: Optional[bytes] = None

    input: Optional[RitualVector] = None

    output: Optional[RitualVector] = None

    flops: Optional[float] = None

    ezkl_proof_path: Optional[Path]
    witness_path: Optional[Path]
    settings_path: Optional[Path]
    vk_path: Optional[Path] = None

    @classmethod
    def from_proof_and_verify_key(
        cls,
        proof: Dict[str, Any],
        verify_key: bytes | str,
        proof_directory: Path = DEFAULT_PROOFS_DIRECTORY,
    ) -> "EZKLProof":
        """
        Create an EZKLProof object from the proof and verification key.

        Args:
            proof (Dict[str, Any]): Proof dictionary
            verify_key (bytes | str): Verification key
            proof_directory (Path, optional): Directory to store the proof.
                Defaults to ~/.cache/ritual/proofs.

        Returns:
            EZKLProof: EZKLProof object
        """
        verify_key = (
            bytes.fromhex(verify_key.removeprefix("0x"))
            if isinstance(verify_key, str)
            else verify_key
        )
        proof_dir, proof_path = get_proof_dir_and_path(proof, proof_directory)
        return cls(
            ezkl_proof=proof,
            verify_key=verify_key,
            witness={},
            ezkl_proof_path=proof_path,
            witness_path=None,
            settings_path=None,
        )

    @classmethod
    def from_files(
        cls,
        proof_path: Path,
        witness_path: Path,
        settings_path: Path,
        flops: Optional[float] = None,
        vk_path: Optional[Path] = None,
    ) -> "EZKLProof":
        """
        Create an EZKLProof object from the files generated by EZKL.

        Args:
            proof_path (Path): Path to the proof file
            witness_path (Path): Path to the witness file
            settings_path (Path): Path to the settings file
            flops (Optional[float], optional): Number of FLOPS in the model.
            vk_path (Optional[Path], optional): Path to the verification key file.
                Defaults to None.

        Returns:
            EZKLProof: EZKLProof object
        """
        with open(proof_path, "r") as f:
            proof = json.loads(f.read())
            public_inputs = proof.get("pretty_public_inputs", {})
            inputs = public_inputs.get("rescaled_inputs")
            inputs = numpy.array(inputs, dtype=float) if inputs else None
            outputs = public_inputs.get("rescaled_outputs")
            outputs = numpy.array(outputs, dtype=float) if outputs else None

        with open(witness_path, "r") as f:
            witness = json.load(f)
            _input, _output = extract_io_from_proof_execution(settings_path, witness)
        if vk_path is not None:
            with open(vk_path, "rb") as f:
                vk = f.read()
        else:
            vk = None

        return cls(
            ezkl_proof=proof,
            verify_key=vk,
            witness=witness,
            input=RitualVector.from_numpy(inputs),
            output=RitualVector.from_numpy(outputs),
            flops=flops,
            ezkl_proof_path=proof_path,
            witness_path=witness_path,
            vk_path=vk_path,
            settings_path=settings_path,
        )


def get_proof_dir_and_path(
    proof: Dict[str, Any],
    proofs_directory: Path = DEFAULT_PROOFS_DIRECTORY,
) -> tuple[Path, Path]:
    """
    Get the proof directory and path for the proof. Prune old items in the proof
    directory to avoid filling up the disk.

    Args:
        proof (Dict[str, Any]): Proof dictionary
        proofs_directory (Path, optional): Directory to store the proof.
            Defaults to ~/.cache/ritual/proofs.

    Returns:
        tuple[Path, Path]: Proof directory and path
    """
    proof_dir = get_proof_dir(proofs_directory)
    proof_path = proof_dir / "proof.json"
    proof_path.write_text(json.dumps(proof))
    return proof_dir, proof_path


async def generate_proof_from_repo_id(
    repo_id: str | RitualRepoId,
    input_vector: np.ndarray[Any, Any],
    hf_token: Optional[str] = None,
) -> EZKLProof:
    """
    Generate a proof for the model using the provided repo_id.

    Args:
        repo_id (str | RitualRepoId): Repo ID for the model
        input_vector (FloatNumpy): Input vector for the model
        hf_token (Optional[str]): Hugging Face API token, optional

    Returns:
        EZKLProof: Proof and verification
    """
    artifacts: EZKLArtifact = RitualArtifactManager.from_repo(
        EZKLArtifact, repo_id, hf_token=hf_token
    ).artifact

    names = get_onnx_input_names(artifacts.onnx_path)
    first = names[0]

    return await generate_proof(
        artifact_files=artifacts,
        prover_key=artifacts.prover_key_path.read_bytes(),
        input_data={first: input_vector},
    )


CACHE_SIZE = 100


@lru_cache(maxsize=CACHE_SIZE)
def get_flops_from_ezkl_aritfact(repo_id: str | RitualRepoId) -> float:
    return get_onnx_flops(
        RitualArtifactManager[EZKLArtifact]
        .from_repo(EZKLArtifact, repo_id)
        .artifact.onnx_path
    )


async def generate_proof(
    artifact_files: EZKLArtifact,
    input_data: ONNXInput,
    output_data: Optional[np.ndarray] = None,  # type: ignore
    prover_key: Optional[bytes] = None,
    proof_directory: Path = DEFAULT_PROOFS_DIRECTORY,
) -> EZKLProof:
    """
    Generate a proof for the model. If a verifier

    Args:
        artifact_files (EZKLArtifact): Paths to the various artifacts generated by
            EZKL
        input_data (ONNXInput): Input data for the model
        output_data (Optional[np.ndarray], optional): Output data for the model.
            Defaults to None.
        prover_key (Optional[bytes], optional): Provide a prover key to generate the
            proof. Defaults to None. If it is not provided, a new circuit setup will
            happen & a new prover key will be generated.
        proof_directory (Optional[Path], optional): Directory to store the proof.
            Defaults to ~/.cache/ritual/proofs.

    Returns:
        EZKLProof: Proof and verification key
    """
    settings_path = artifact_files.settings_path
    compiled_model_path = artifact_files.compiled_model_path
    srs_path = artifact_files.srs_path
    vk_path = artifact_files.verifier_key_path
    pk_path = settings_path.parent / "test.pk"
    prover_key = prover_key or artifact_files.prover_key_path.read_bytes()

    if srs_path is None:
        srs_path = settings_path.parent / "srs_file.srs"
        res = await ezkl.get_srs(settings_path, srs_path=srs_path)
        assert res is True

    if prover_key is None:
        # if pk is not provided, then we'll generate both pk & vk
        log.info("proving key not provided, performing a new ZK setup")
        vk_path = settings_path.parent / "test.vk"
        res = ezkl.setup(
            compiled_model_path, vk_path=vk_path, pk_path=pk_path, srs_path=srs_path
        )
        assert res is True
        assert os.path.isfile(vk_path)
        assert os.path.isfile(pk_path)
    else:
        log.info("proving key provided using the existing setup")
        pk_path.write_bytes(prover_key)

    proof_dir, proof_path = get_proof_dir_and_path({}, proof_directory)

    os.makedirs(proof_dir, exist_ok=True)

    witness_path = proof_dir / "witness.json"
    data_path = proof_dir / "data.json"

    input_data_ = list(input_data.values())[0].reshape([-1]).tolist()

    data = dict(input_data=[input_data_])

    if output_data is not None:
        data["output_data"] = [output_data.reshape([-1]).tolist()]

    json.dump(data, open(data_path, "w"))

    await ezkl.gen_witness(
        data_path, model=compiled_model_path, output=witness_path, srs_path=srs_path
    )

    assert os.path.isfile(witness_path)

    ONNXModelAnalyzer(artifact_files.onnx_path).calculate_flops()

    res = ezkl.prove(
        witness=witness_path,
        model=compiled_model_path,
        pk_path=pk_path,
        proof_path=proof_path,
        proof_type="single",
        srs_path=srs_path,
    )

    assert res, "unable to generate proof"

    assert os.path.isfile(proof_path)

    if vk_path:
        log.debug("vk available, verifying the proof")
        verify_success = ezkl.verify(
            proof_path=proof_path,
            settings_path=settings_path,
            vk_path=vk_path,
            srs_path=srs_path,
        )
        assert verify_success, "unable to verify generated proof"

    if artifact_files.onnx_path:
        flops = get_onnx_flops(artifact_files.onnx_path)
    else:
        flops = None

    return EZKLProof.from_files(
        proof_path=proof_path,
        witness_path=witness_path,
        settings_path=settings_path,
        flops=flops,
        vk_path=vk_path,
    )


async def verify_proof_from_repo(
    proof_dict: Dict[str, Any], repo_id: str | RitualRepoId
) -> bool:
    files: EZKLArtifact = (
        RitualArtifactManager[EZKLArtifact].from_repo(EZKLArtifact, repo_id).artifact
    )
    assert files.verifier_key_path.exists()
    proof = EZKLProof.from_proof_and_verify_key(
        proof=proof_dict, verify_key=files.verifier_key_path.read_bytes()
    )
    return await verify_proof(proof, files)


async def verify_proof(proof: EZKLProof, artifact_files: EZKLArtifact) -> bool:
    """
    Verify a proof using the provided artifacts.

    Args:
        proof (EZKLProof): Proof to verify
        artifact_files (EZKLArtifact): Paths to the various artifacts generated
        by EZKL

    Returns:
        bool: True if the proof is valid, False otherwise
    """

    with tempfile.TemporaryDirectory() as _dir:
        temp_dir = Path(_dir)
        vk_path = temp_dir / "test.vk"
        assert proof.verify_key
        vk_path.write_bytes(proof.verify_key)

        settings_path = artifact_files.settings_path
        if not artifact_files.srs_path or not artifact_files.srs_path.exists():
            res = await ezkl.get_srs(settings_path, srs_path=artifact_files.srs_path)
            assert res is True

        verify_success = ezkl.verify(
            proof_path=proof.ezkl_proof_path,
            settings_path=settings_path,
            vk_path=vk_path,
            srs_path=artifact_files.srs_path,
        )

        return bool(verify_success)
