"""LLM execution utilities backed by the Ember runtime."""

import json
import logging
import random
import re
import threading
import time
from beartype.typing import Any, Dict, List, Optional, Tuple
from collections import deque

from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential_jitter,
)

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Local token counting (no API calls)
# ---------------------------------------------------------------------------

try:
    import tiktoken as _tiktoken

    _CL100K = _tiktoken.get_encoding("cl100k_base")
except Exception:  # pragma: no cover – tiktoken optional at import time
    _CL100K = None


def _count_tokens_local(text: str) -> int:
    """Estimate token count locally without making any API calls.

    Uses tiktoken cl100k_base when available (covers OpenAI family), otherwise
    falls back to a ~4 chars/token heuristic.
    """
    if not text:
        return 0
    if _CL100K is not None:
        return len(_CL100K.encode(text))
    return max(1, len(text) // 4)


class AdaptiveCongestionController:
    """Adaptive congestion control based on error rates.

    Implements TCP-like AIMD (Additive Increase, Multiplicative Decrease):
    - On success: gradually increase throughput
    - On rate limit errors: multiplicatively decrease and back off

    This provides smooth adaptation to API rate limits without knowing
    the exact limits in advance.
    """

    def __init__(
        self,
        window_size: int = 20,
        error_threshold: float = 0.25,
        base_delay: float = 0.1,
        max_delay: float = 30.0,
    ):
        """Initialize congestion controller.

        Args:
            window_size: Number of recent requests to track for error rate.
            error_threshold: Error rate above which to back off (0.25 = 25%).
            base_delay: Minimum delay between requests (seconds).
            max_delay: Maximum delay cap (seconds).
        """
        self._lock = threading.Lock()
        self._window: deque[bool] = deque(maxlen=window_size)
        self._current_delay = base_delay
        self._base_delay = base_delay
        self._max_delay = max_delay
        self._error_threshold = error_threshold
        self._last_request_time = 0.0

    def record_success(self) -> None:
        """Record a successful request. Gradually reduce delay."""
        with self._lock:
            self._window.append(True)
            # Additive increase: slowly reduce delay on success
            self._current_delay = max(
                self._base_delay, self._current_delay * 0.95
            )

    def record_rate_limit(self, retry_after: float | None = None) -> None:
        """Record a rate limit error. Increase delay significantly.

        Args:
            retry_after: If the API provided a Retry-After value, use it.
        """
        with self._lock:
            self._window.append(False)
            if retry_after:
                # Use the API's suggested delay
                self._current_delay = min(retry_after, self._max_delay)
            else:
                # Multiplicative decrease: double the delay on rate limit
                self._current_delay = min(
                    self._current_delay * 2.0, self._max_delay
                )

    def record_error(self) -> None:
        """Record a non-rate-limit error. Small delay increase."""
        with self._lock:
            self._window.append(False)
            # Smaller increase for general errors
            self._current_delay = min(
                self._current_delay * 1.2, self._max_delay
            )

    def get_error_rate(self) -> float:
        """Get current error rate over the window."""
        with self._lock:
            if not self._window:
                return 0.0
            return 1.0 - (sum(self._window) / len(self._window))

    def wait_if_needed(self) -> None:
        """Wait if error rate is high or delay is set."""
        with self._lock:
            delay = self._current_delay
            error_rate = self.get_error_rate()

        # Extra delay if error rate is high
        if error_rate > self._error_threshold:
            delay = delay * (1 + error_rate)

        if delay > 0:
            time.sleep(delay)


# Global congestion controller shared across all LLM calls
_congestion_controller = AdaptiveCongestionController()


def _extract_retry_after(error_str: str) -> float | None:
    """Extract retry-after value from error message if present.

    OpenAI often includes "Please retry after X seconds" or similar in error messages.
    """
    patterns = [
        r"retry.?after[:\s]+(\d+(?:\.\d+)?)",  # "retry after: 30" or "retry-after 30"
        r"(\d+(?:\.\d+)?)\s*seconds?",          # "30 seconds"
        r"wait[:\s]+(\d+(?:\.\d+)?)",           # "wait: 30" or "wait 30"
    ]
    for pattern in patterns:
        match = re.search(pattern, error_str, re.IGNORECASE)
        if match:
            try:
                return float(match.group(1))
            except ValueError:
                pass
    return None


# Handle both package import and direct script execution
try:
    from .utils import (
        count_tokens,
        evaluate_bbeh,
        exact_match_score,
        extract_mc_answer as extract_mc_answer_utils,
        extract_yes_no_answer,
        f1_score,
        get_bert_score,
        model_prompting,
    )
except ImportError:
    from utils import (
        count_tokens,
        evaluate_bbeh,
        exact_match_score,
        extract_mc_answer as extract_mc_answer_utils,
        extract_yes_no_answer,
        f1_score,
        get_bert_score,
        model_prompting,
    )


def extract_numeric_answer(text: str) -> Optional[str]:
    """Extract a numeric answer from an LLM response to a math problem.

    Uses the official GSM8K evaluation approach: look for explicit answer
    markers first, then fall back to the last number in the response.

    Priority order (high-confidence markers first):
      1. \\boxed{X}  — LaTeX convention
      2. <X>         — angle-bracket convention
      3. #### X      — GSM8K ground-truth format
      4. "answer is/=/: X" or "final answer: X"
      5. Last **X** bold number in response (common LLM convention)
      6. Last standalone number (official GSM8K fallback)

    The `= X` pattern is intentionally omitted: `re.search` returns the
    first match, which in chain-of-thought responses is an intermediate
    computation step, not the final answer.  The "last number" fallback
    handles this correctly by taking the final occurrence.

    Args:
        text: The LLM response text.

    Returns:
        The extracted numeric string, or None if no number found.
    """
    if not text:
        return None

    # Normalize: remove commas from numbers (e.g., "1,234" -> "1234")
    text_normalized = re.sub(r"(\d),(\d)", r"\1\2", text)

    # Pattern for a number: optional minus, digits, optional decimal part
    num_pattern = r"-?\d+(?:\.\d+)?"

    # Priority 1: Explicit, unambiguous answer markers.
    # Use findall + take LAST match (not re.search which returns FIRST).
    # In chain-of-thought, LLMs may write "answer is 24" then self-correct
    # to "answer is 42" — the last occurrence is the final answer.
    #
    # The negative lookahead (?!\s*\.\s*[A-Za-z]) on answer_is patterns
    # prevents matching numbered-list items: "answer:\n\n1. First step"
    # would otherwise extract "1" as the answer.
    explicit_patterns = [
        rf"\\boxed{{({num_pattern})}}",                              # \boxed{42}
        rf"\<({num_pattern})\>",                                     # <42>
        rf"####\s*({num_pattern})",                                  # #### 42
        rf"(?:final\s+)?answer\s*(?:is|=|:)\s*\$?({num_pattern})(?!\s*\.\s*[A-Za-z])",
        rf"(?:result|solution)\s*(?:is|=|:)\s*\$?({num_pattern})(?!\s*\.\s*[A-Za-z])",
    ]
    for pattern in explicit_patterns:
        matches = re.findall(
            pattern, text_normalized, re.IGNORECASE | re.MULTILINE,
        )
        if matches:
            return matches[-1]

    # Priority 2: Bold answer extraction — LAST bold block only.
    # LLMs bold their final answers in two styles:
    #   a) Number inside bold: **42**, **$5.99**, **240 + 160 = 400**
    #   b) Label in bold, number outside: **Answer:** 10
    # Only the LAST bold block is checked.  If it contains numbers,
    # return the last one (handles equations).  If it contains no
    # numbers (like "**Answer:**"), skip bold entirely so Priority 1
    # (answer_is patterns) or Priority 3 (last standalone number) can
    # extract the number from the surrounding plain text.
    bold_blocks = re.findall(
        r"\*\*([^*\n]+?)\*\*", text_normalized,
    )
    if bold_blocks:
        last_block = bold_blocks[-1]
        last_block_nums = re.findall(
            num_pattern, last_block.replace(",", ""),
        )
        if last_block_nums:
            # Equations ("240 + 160 = 400"): answer is the LAST number.
            # Sentences ("96 people in 2 days"): answer is the FIRST.
            if "=" in last_block:
                return last_block_nums[-1]
            return last_block_nums[0]

    # Priority 3: Last standalone number in the text.
    # This is the official GSM8K fallback: the final number in the
    # chain-of-thought is almost always the answer.
    # The lookahead (?!\d) prevents matching a partial number (e.g.,
    # "3" from "3.14"), but allows sentence-ending periods ("42.").
    # num_pattern already handles decimals ("3.14" is one token).
    last_number = re.findall(
        rf"(?<![.\d])({num_pattern})(?!\d)", text_normalized,
    )
    if last_number:
        return last_number[-1]

    return None


def _normalize_cruxeval_output(text: str) -> str:
    """Normalize a CRUXEval output prediction for comparison.

    Handles two systematic issues:
      1. Verbose chain-of-thought: extracts the final answer from
         explanations like "Let me trace through... The output is X".
      2. Quoting mismatch: ground truths use Python repr (e.g., 'foo')
         but models may output bare strings (foo) or double-quoted ("foo").

    The official CRUXEval evaluation uses execution-based comparison
    (assert f(input) == output), which handles formatting automatically.
    This function approximates that for cached text responses.
    """
    if not text:
        return ""
    text = text.strip()

    # If response is short (likely just the answer), normalize directly.
    if len(text) < 200:
        return _strip_outer_quotes(text)

    # Verbose response: extract the final answer.
    # Look for common answer markers (last occurrence wins).
    # Priority: code blocks > inline backticks > prose markers.
    # Code blocks are the most deliberate answer format (models wrap
    # final output in ```) so they must be checked before inline
    # backtick patterns that could match explanation fragments.
    answer_patterns = [
        r"```\s*\n?([^`]+?)\n?```\s*$",
        r"(?:output|result|returns?|produces?|yields?|evaluates? to|is)\s*(?:is|=|:)\s*`([^`]+)`",
        r"(?:output|result|returns?)\s*(?:is|=|:)\s*(.+?)(?:\n|$)",
    ]
    for pattern in answer_patterns:
        matches = re.findall(pattern, text, re.IGNORECASE | re.MULTILINE)
        if matches:
            return _strip_outer_quotes(matches[-1].strip())

    # Fallback: last non-empty line (often the final answer).
    lines = [ln.strip() for ln in text.strip().splitlines() if ln.strip()]
    if lines:
        last = lines[-1]
        # Skip lines that look like prose continuation
        if not last.startswith(("Let me", "So ", "Now ", "Therefore",
                                "The function", "This ", "I ")):
            return _strip_outer_quotes(last)

    return _strip_outer_quotes(text)


def _strip_outer_quotes(s: str) -> str:
    """Remove one layer of matching quotes from a string value."""
    s = s.strip()
    if len(s) >= 2:
        if (s[0] == "'" and s[-1] == "'") or (s[0] == '"' and s[-1] == '"'):
            return s[1:-1]
    return s


def _cruxeval_literal_equal(pred: str, gt: str) -> bool:
    """Compare CRUXEval outputs as Python objects.

    Replicates the official execution-based evaluation semantics
    (``assert output == prediction``) using ``ast.literal_eval``,
    which is safe (only parses literals, no code execution).

    Handles cases that text comparison misses:
      - ``True == 1``, ``False == 0`` (Python equality)
      - ``[1,2]`` vs ``[1, 2]`` (whitespace in containers)
      - ``'hello'`` vs ``"hello"`` (quote style)

    Falls back to text comparison when parsing fails (e.g. the
    model produced an explanation instead of a Python literal).
    """
    import ast

    # Try structural comparison first.
    try:
        pred_obj = ast.literal_eval(pred)
        gt_obj = ast.literal_eval(gt)
        return pred_obj == gt_obj
    except (ValueError, SyntaxError, TypeError):
        # TypeError: unhashable types in containers (e.g. dict inside set).
        pass

    # Fall back to exact text match (after normalization already
    # applied by the caller via _normalize_cruxeval_output).
    return pred == gt


class LLMEngine:
    """
    A class to manage interactions with multiple language models and evaluate their performance.

    Handles model selection, querying, cost calculation, and performance evaluation
    using various metrics for different tasks.
    """

    def __init__(self, llm_names: List[str], llm_description: Dict[str, Dict[str, Any]]):
        """
        Initialize the LLM Engine with available models and their descriptions.

        Args:
            llm_names: List of language model names available in the engine
            llm_description: Dictionary containing model configurations and pricing details
                Structure: {
                    "model_name": {
                        "model": "api_identifier",
                        "input_price": cost_per_input_token,
                        "output_price": cost_per_output_token,
                        ...
                    },
                    ...
                }
        """
        self.llm_names = llm_names
        self.llm_description = llm_description

    def compute_cost(self, llm_idx: int, input_text: str, output_text: str) -> float:
        """Calculate cost using local tiktoken estimation.

        Uses cl100k_base as a universal tokenizer approximation across all
        providers.  Token counts may differ from actual provider billing by
        ~10-20%, but that is more than sufficient for the GNN's relative cost
        signal (min-max normalized per task).

        Args:
            llm_idx: Index of the model in the llm_names list.
            input_text: The input prompt sent to the model.
            output_text: The actual response text from the model.

        Returns:
            Estimated USD cost (per-million pricing converted to per-call).
        """
        llm_name = self.llm_names[llm_idx]

        # Local token counting via tiktoken — no API calls.
        input_size = _count_tokens_local(input_text)
        output_size = _count_tokens_local(output_text)

        input_price = self.llm_description[llm_name]["input_price"]
        output_price = self.llm_description[llm_name]["output_price"]

        return ((input_size / 1_000_000) * input_price) + (
            (output_size / 1_000_000) * output_price
        )

    def get_llm_response(self, query: str, llm_idx: int) -> str:
        """Send a query to a language model and get its response.

        No artificial output-token limits are imposed.  Each provider uses
        its own model-maximum default; models stop at natural completion.

        Args:
            query: The prompt text to send to the model.
            llm_idx: Index of the model in the llm_names list.

        Returns:
            The model's text response.
        """
        llm_name = self.llm_names[llm_idx]
        model = self.llm_description[llm_name]["model"]
        return self._call_with_backoff(model, query, llm_name)

    @retry(
        stop=stop_after_attempt(6),
        wait=wait_exponential_jitter(initial=1, max=60, jitter=5),
        reraise=True,
    )
    def _call_with_backoff(
        self, model: str, query: str, llm_name: str
    ) -> str:
        """Execute API call with exponential backoff and jitter."""
        try:
            logger.info("→ %s starting...", llm_name)
            start = time.time()
            response = model_prompting(llm_model=model, prompt=query)
            elapsed = time.time() - start

            if not response or not response.strip():
                logger.warning(
                    "⚠ %s returned empty response after %.1fs",
                    llm_name,
                    elapsed,
                )

            logger.info("← %s done in %.1fs", llm_name, elapsed)
            return response
        except Exception as e:
            logger.warning("✗ %s failed: %s", llm_name, e)
            raise


    def get_llm_response_with_latency(self, query: str, llm_idx: int) -> Tuple[str, float]:
        """
        Execute a query and return both the response and wall-clock latency in seconds.
        """
        start = time.perf_counter()
        response = self.get_llm_response(query=query, llm_idx=llm_idx)
        elapsed = time.perf_counter() - start
        return response, float(elapsed)

    def eval(self, prediction: str, ground_truth: str, metric: str) -> float:
        """Evaluate the model's prediction against the ground truth.

        Args:
            prediction: The model's output text.
            ground_truth: The correct expected answer.
            metric: The evaluation metric to use. Supported:
                - 'em': Exact match after normalization
                - 'em_mc': Multiple choice exact match (extracts answer letter)
                - 'bert_score': BERT-based semantic similarity
                - 'GSM8K': Math problem evaluation (extracts numeric answer)
                - 'f1_score': Token-level F1 for QA tasks

        Returns:
            Evaluation score between 0 and 1.
        """
        # Exact match evaluation
        if metric == "em":
            result = exact_match_score(prediction, ground_truth)
            return float(result)

        # CRUXEval output prediction: compare as Python objects using
        # ast.literal_eval, which replicates the official execution-based
        # evaluation (assert output == prediction) without running code.
        # This correctly handles True==1, whitespace in containers,
        # and quote-style differences that text comparison misses.
        elif metric == "cruxeval_em":
            pred_n = _normalize_cruxeval_output(prediction)
            gt_n = _strip_outer_quotes(ground_truth.strip())
            if _cruxeval_literal_equal(pred_n, gt_n):
                return 1.0
            return 0.0

        # Multiple choice exact match - uses consolidated answer extraction from utils
        elif metric == "em_mc":
            pred_answer = extract_mc_answer_utils(prediction)
            gt_answer = extract_mc_answer_utils(ground_truth)

            # If ground truth is already a single letter, use it directly
            if len(ground_truth.strip()) == 1 and ground_truth.strip().isalpha():
                gt_answer = ground_truth.strip().upper()

            # Return 0.0 if no valid answer extracted from either side
            # Critical: prevents None == None from being treated as a match
            if not pred_answer or not gt_answer:
                return 0.0

            return 1.0 if pred_answer == gt_answer else 0.0

        # BERT-based semantic similarity score
        elif metric == "bert_score":
            result = get_bert_score([prediction], [ground_truth])
            return result

        # GSM8K math problem evaluation - uses improved answer extraction
        elif metric == "GSM8K":
            # Extract the final answer from ground truth (after the "####" delimiter)
            gt_clean = ground_truth.split("####")[-1].strip()
            # Remove any commas from ground truth number
            gt_clean = gt_clean.replace(",", "")

            # Use improved extraction for prediction
            pred_answer = extract_numeric_answer(prediction)

            if pred_answer is None:
                return 0.0

            # Numeric comparison handles decimal format mismatch:
            # "990.00" vs "990", "42.0" vs "42", "-3.00" vs "-3"
            try:
                if float(pred_answer) == float(gt_clean):
                    return 1.0
            except ValueError:
                pass

            return 1.0 if pred_answer == gt_clean else 0.0

        # SQuAD evaluation: containment with F1 fallback.
        #
        # Primary: check if normalized GT appears in normalized response
        # (handles verbose LLM outputs like "The answer is X").
        # Fallback: token-level F1 gives partial credit for paraphrased
        # answers (e.g., "nationally" vs "in the nation") instead of 0.
        #
        # normalize_answer handles Unicode (NFKD decomposition) so smart
        # quotes, en dashes, etc. are stripped consistently.
        elif metric == "squad_em":
            try:
                from .utils import normalize_answer as _norm
                from .utils import f1_score as _f1
            except ImportError:
                from utils import normalize_answer as _norm
                from utils import f1_score as _f1
            pred_n = _norm(prediction) or ""
            gt_n = _norm(ground_truth) or ""
            if not gt_n:
                return 0.0
            if gt_n in pred_n:
                return 1.0
            # Fallback: token F1 for paraphrased answers.
            f1, _, _ = _f1(prediction, ground_truth)
            return float(f1)

        # Binary classification exact match — for tasks where the ground
        # truth is strictly "0" or "1" (e.g., codexglue_defect: "is this
        # code defective?").  Some models wrap the label in explanations:
        #   "1. The code is defective because..."
        #   "The answer is 0 (not defective)."
        # We extract the first standalone 0 or 1 from the response.
        elif metric == "binary_em":
            gt_clean = ground_truth.strip()
            if gt_clean not in ("0", "1"):
                return 0.0
            pred_text = prediction.strip()
            # Short response: exact match after stripping whitespace
            if pred_text in ("0", "1"):
                return 1.0 if pred_text == gt_clean else 0.0
            # Verbose response: extract first standalone 0 or 1.
            # Require word boundary to avoid matching digits within
            # larger numbers like "100" or variable names.
            match = re.search(r"(?<!\d)\b([01])\b(?!\d)", pred_text)
            if match:
                return 1.0 if match.group(1) == gt_clean else 0.0
            return 0.0

        # F1 score for partial matching (used in QA tasks)
        elif metric == "f1_score":
            f1, prec, recall = f1_score(prediction, ground_truth)
            return f1

        # Exact numeric match - for BBEH arithmetic, buggy tables, etc.
        elif metric == "em_numeric":
            # Normalize ground truth: strip whitespace, remove commas
            gt_clean = ground_truth.strip().replace(",", "")

            # Extract numeric answer from prediction
            pred_answer = extract_numeric_answer(prediction)

            if pred_answer is None:
                return 0.0

            # Handle decimal comparison (e.g., "404.0" vs "404")
            try:
                pred_float = float(pred_answer)
                gt_float = float(gt_clean)
                # Use exact float comparison for integers, tolerance for decimals
                if pred_float == gt_float:
                    return 1.0
                # Allow small floating point tolerance for decimal answers
                if abs(pred_float - gt_float) < 1e-6:
                    return 1.0
            except ValueError:
                pass

            # Fall back to string comparison if not valid floats
            return 1.0 if pred_answer == gt_clean else 0.0

        # Yes/No/Ambiguous exact match - extracts answer from verbose responses
        elif metric == "yes_no":
            pred_answer = extract_yes_no_answer(prediction)
            gt_clean = ground_truth.strip()

            # If no answer extracted from prediction, it's wrong
            if pred_answer is None:
                return 0.0

            # Case-insensitive comparison
            return 1.0 if pred_answer.lower() == gt_clean.lower() else 0.0

        # BBEH (BIG-Bench Extra Hard) evaluation using official methodology
        # Extracts final answer after "The answer is:" prefix and uses fuzzy matching
        elif metric == "bbeh":
            return evaluate_bbeh(prediction, ground_truth)

        # OfficeQA numeric: 1% relative tolerance for financial data.
        # Handles verbose LLM responses by extracting numeric answer first.
        # Non-numeric ground truths fall back to containment check.
        elif metric == "officeqa_numeric":
            gt_str = ground_truth.strip().replace(",", "")
            try:
                gt_num = float(gt_str)
            except ValueError:
                # Non-numeric ground truth: containment check.
                try:
                    from .utils import normalize_answer as _norm
                except ImportError:
                    from utils import normalize_answer as _norm
                pred_n = _norm(prediction) or ""
                gt_n = _norm(ground_truth) or ""
                return 1.0 if gt_n and gt_n in pred_n else 0.0

            pred_answer = extract_numeric_answer(prediction)
            if pred_answer is None:
                return 0.0
            try:
                pred_float = float(pred_answer.replace(",", ""))
            except ValueError:
                return 0.0
            if abs(gt_num) < 1e-10:
                return 1.0 if abs(pred_float) < 1e-10 else 0.0
            return 1.0 if abs(pred_float - gt_num) / abs(gt_num) < 0.01 else 0.0

        # Default case for unrecognized metrics
        else:
            return 0.0
