"""Utility helpers for dataset construction in embergraphrouter."""

import json
import os
import pickle
from typing import List, Optional, Tuple

import numpy as np
import torch
from bert_score import score
from sentence_transformers import SentenceTransformer
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

# Ember imports are optional - only needed for remote embedding/LLM calls
try:
    from ember._internal.context.runtime import EmberContext
    from ember.models.runtime.registry import ModelRegistry
    from ember.models.tokenization import count_tokens
    _EMBER_AVAILABLE = True
except ImportError:
    _EMBER_AVAILABLE = False
    EmberContext = None  # type: ignore[misc, assignment]
    ModelRegistry = None  # type: ignore[misc, assignment]
    count_tokens = None  # type: ignore[assignment]

_ENCODER_NAME = os.environ.get("EGR_TEXT_ENCODER", "google/gemini-embedding-001")
_ENCODER_DIM = int(os.environ.get("EGR_EMBEDDING_DIM", "256"))
_SUPPORTED_REMOTE_PROVIDERS = frozenset({"openai", "google"})
_REMOTE_ENCODER_ID: Optional[str]
_local_encoder: Optional[SentenceTransformer]
if "/" in _ENCODER_NAME:
    provider, _, model_name = _ENCODER_NAME.partition("/")
    if provider not in _SUPPORTED_REMOTE_PROVIDERS or not model_name:
        raise ValueError(
            f"EGR_TEXT_ENCODER must be a SentenceTransformer name or "
            f"'<provider>/<model>' where provider is one of: {sorted(_SUPPORTED_REMOTE_PROVIDERS)}"
        )
    _REMOTE_ENCODER_ID = _ENCODER_NAME
    _local_encoder = None
else:
    _REMOTE_ENCODER_ID = None
    _local_encoder = SentenceTransformer(_ENCODER_NAME)

# Registry requires ember - set to None if unavailable
_registry = ModelRegistry(context=EmberContext.current()) if _EMBER_AVAILABLE else None


# File I/O functions
def loadjson(filename: str) -> dict:
    """
    Load data from a JSON file.

    Args:
        filename: Path to the JSON file

    Returns:
        Dictionary containing the loaded JSON data
    """
    with open(filename, "r", encoding="utf-8") as file:
        data = json.load(file)
    return data


def savejson(data: dict, filename: str) -> None:
    """
    Save data to a JSON file.

    Args:
        data: Dictionary to save
        filename: Path where the JSON file will be saved
    """
    with open(filename, "w") as json_file:
        json.dump(data, json_file, indent=4)


def loadpkl(filename: str) -> any:
    """
    Load data from a pickle file.

    Args:
        filename: Path to the pickle file

    Returns:
        The unpickled object
    """
    with open(filename, "rb") as file:
        data = pickle.load(file)
    return data


def savepkl(data: any, filename: str) -> None:
    """
    Save data to a pickle file.

    Args:
        data: Object to save
        filename: Path where the pickle file will be saved
    """
    with open(filename, "wb") as pkl_file:
        pickle.dump(data, pkl_file)


# Text normalization and evaluation functions
#
# Canonical implementations live in ember.utils.eval.  These functions
# delegate via lazy imports so that embergraphrouter callers (which may
# pull in torch/numpy) don't force ember.utils.eval at module load time.


def extract_mc_answer(text: str) -> Optional[str]:
    """Extract multiple choice answer letter from text.

    Delegates to :func:`ember.utils.eval.extractors.extract_mc_answer`.
    """
    from ember.utils.eval.extractors import extract_mc_answer as _impl

    return _impl(text)


def extract_yes_no_answer(text: str) -> Optional[str]:
    """Extract Yes/No/Ambiguous answer from LLM response text.

    Delegates to :func:`ember.utils.eval.extractors.extract_yes_no_answer`.
    """
    from ember.utils.eval.extractors import extract_yes_no_answer as _impl

    return _impl(text)


def normalize_answer(s: str, normal_method: str = "") -> Optional[str]:
    """Normalize text for evaluation.

    For ``normal_method="mc"``, delegates to :func:`extract_mc_answer`.
    Otherwise delegates to :func:`ember.utils.eval.evaluators.normalize_answer`.

    Args:
        s: String to normalize (non-string inputs coerced to str).
        normal_method: ``"mc"`` for multiple choice, ``""`` for standard.

    Returns:
        Normalized string, or ``None`` for MC mode if no valid answer found.
    """
    if not isinstance(s, str):
        s = str(s)
    if normal_method == "mc":
        return extract_mc_answer(s)
    from ember.utils.eval.evaluators import normalize_answer as _impl

    return _impl(s)


def f1_score(prediction: str, ground_truth: str) -> Tuple[float, float, float]:
    """Token-level F1 between prediction and ground truth.

    Delegates to :func:`ember.utils.eval.evaluators.token_f1_score`.
    """
    from ember.utils.eval.evaluators import token_f1_score

    return token_f1_score(prediction, ground_truth)


def exact_match_score(
    prediction: str, ground_truth: str, normal_method: str = "",
) -> bool:
    """Check if prediction matches ground truth after normalization.

    For ``normal_method="mc"``, uses MC extraction (returns ``False``
    if either side has no valid MC answer).  Otherwise delegates to
    :func:`ember.utils.eval.evaluators.exact_match_score`.
    """
    if normal_method == "mc":
        pred_normalized = normalize_answer(prediction, normal_method="mc")
        gt_normalized = normalize_answer(ground_truth, normal_method="mc")
        if pred_normalized is None or gt_normalized is None:
            return False
        return pred_normalized == gt_normalized
    from ember.utils.eval.evaluators import exact_match_score as _impl

    return _impl(prediction, ground_truth)


# BBEH (BIG-Bench Extra Hard) evaluation functions
# Based on official evaluation from https://github.com/google-deepmind/bbeh


def strip_latex_bbeh(response: str) -> str:
    """Strip LaTeX formatting from BBEH answers.

    Delegates to :func:`ember.utils.eval.evaluators.strip_latex_bbeh`.
    """
    from ember.utils.eval.evaluators import strip_latex_bbeh as _impl

    return _impl(response)


def extract_answer_bbeh(sample: str) -> str:
    """Extract final answer from BBEH model response.

    Delegates to :func:`ember.utils.eval.evaluators.extract_answer_bbeh`.
    """
    from ember.utils.eval.evaluators import extract_answer_bbeh as _impl

    return _impl(sample)


def fuzzy_match_bbeh(prediction: str, reference: str) -> bool:
    """Fuzzy match for BBEH evaluation.

    Delegates to :func:`ember.utils.eval.evaluators.fuzzy_match_bbeh`.
    """
    from ember.utils.eval.evaluators import fuzzy_match_bbeh as _impl

    return _impl(prediction, reference)


def evaluate_bbeh(prediction: str, ground_truth: str) -> float:
    """Evaluate BBEH task using official methodology with MC fallback.

    Delegates to :func:`ember.utils.eval.evaluators.evaluate_bbeh`.
    """
    from ember.utils.eval.evaluators import evaluate_bbeh as _impl

    return _impl(prediction, ground_truth)


def get_bert_score(generate_response: List[str], ground_truth: List[str]) -> float:
    """
    Calculate BERT score between generated responses and ground truths.

    Args:
        generate_response: List of generated responses
        ground_truth: List of ground truth texts

    Returns:
        Average BERT score (F1)
    """
    F_l = []
    for inter in range(len(generate_response)):
        generation = generate_response[inter]
        gt = ground_truth[inter]
        P, R, F = score([generation], [gt], lang="en", verbose=True)
        F_l.append(F.mean().numpy().reshape(1)[0])
    return np.array(F_l).mean()


# Embedding and dimensionality reduction
def reduce_embedding_dim(embed: np.ndarray, dim: int = 50) -> np.ndarray:
    """
    Reduce dimensionality of embeddings using PCA.

    Args:
        embed: Embedding vectors
        dim: Target dimension

    Returns:
        Reduced embeddings
    """
    pca = PCA(n_components=dim)
    reduced_embeddings = pca.fit_transform(embed)
    return reduced_embeddings


def get_embedding(instructions: List[str], dimensions: Optional[int] = None) -> np.ndarray:
    """Return sentence embeddings for the provided texts.

    Args:
        instructions: List of text strings to embed.
        dimensions: Output dimensionality. For remote providers supporting
            Matryoshka embeddings, this requests truncated vectors directly.
            Defaults to EGR_EMBEDDING_DIM env var (256) for remote providers.
            Ignored for local SentenceTransformer models.

    Returns:
        Array of shape (len(instructions), embedding_dim).
    """
    if _local_encoder is not None:
        return _local_encoder.encode(instructions)

    if _REMOTE_ENCODER_ID is None:
        raise RuntimeError("No embedding backend configured.")

    if _registry is None:
        raise RuntimeError(
            "Remote embedding requires ember package. Install with: pip install ember"
        )
    dim = dimensions if dimensions is not None else _ENCODER_DIM
    response = _registry.embed_model(_REMOTE_ENCODER_ID, instructions, dimensions=dim)
    return np.asarray(response.embeddings, dtype=np.float32)


# LLM prompting
def model_prompting(
    llm_model: str,
    prompt: str,
    return_num: Optional[int] = 1,
    max_token_num: Optional[int] = None,
    temperature: Optional[float] = None,
    top_p: Optional[float] = None,
    stream: Optional[bool] = None,
) -> str:
    """Retrieve a completion using Ember's provider runtime.

    Args:
        llm_model: Canonical model identifier (provider/model).
        prompt: User message supplied to the model.
        return_num: Number of completions requested (currently unused).
        max_token_num: Maximum tokens requested for the generation. If None,
            uses provider defaults which is recommended for reasoning models.
        temperature: Sampling temperature.
        top_p: Nucleus sampling parameter.
        stream: Whether to stream the response (unsupported in this path).

    Returns:
        The generated text response.
    """
    if _registry is None:
        raise RuntimeError(
            "LLM prompting requires ember package. Install with: pip install ember"
        )
    provider = _registry.get_model(llm_model)
    invocation_kwargs: dict[str, object] = {"model": llm_model}
    if max_token_num is not None:
        invocation_kwargs["max_tokens"] = max_token_num
    if temperature is not None:
        invocation_kwargs["temperature"] = temperature
    if top_p is not None:
        invocation_kwargs["top_p"] = top_p
    if stream is not None:
        invocation_kwargs["stream"] = stream

    response = provider.complete(prompt, **invocation_kwargs)
    return response.data
