"""
Response parsers for MCQ answers.

Two styles matched to util.prompts:
  - parse_mcq_json(text, qid)     — strict/lenient JSON extraction
  - parse_mcq_letters(text)       — letter-regex fallback for prose outputs

Both return a single-letter normalized answer string ('A' / 'B' / 'C' / 'D')
or '' on failure.
"""
from __future__ import annotations

import json
import re

from util.benchmark import normalize_answers


_LETTERS = {'A', 'B', 'C', 'D'}


def norm_qid(q) -> str:
    """
    Normalize qid to a bare string. '13', '13.0', 13, 13.0 all become '13'.
    This is in particular useful for pandas for pretty table printing.
    """
    s = str(q).strip()
    if s.endswith('.0'):
        s = s[:-2]
    return s


def parse_json_flexible(text: str):
    """Parse any JSON object/array from text. Falls back to regex around obvious braces."""
    if not text:
        return None
    try:
        return json.loads(text)
    except (json.JSONDecodeError, TypeError):
        pass
    # Find the first plausible JSON array or object
    for pat in (r'\[.*\]', r'\{.*\}'):
        m = re.search(pat, text, re.DOTALL)
        if m:
            try:
                return json.loads(m.group(0))
            except json.JSONDecodeError:
                continue
    return None


_NUMERIC_TO_LETTER = {'1': 'A', '2': 'B', '3': 'C', '4': 'D'}

def _normalize_letter_value(value) -> str:
    """Return the first valid MCQ letter in a model answer value."""
    if isinstance(value, list):
        value = value[0] if value else ''
    text = str(value).strip().upper()
    if text in _LETTERS:
        return text
    # Numeric index to letter mapping. Some models (videochat-flash, smaller
    # local VLs) consistently emit `"answer": "3"` meaning the 3rd option.
    # Treat 1/2/3/4 as A/B/C/D when the value is JUST a digit.
    if text in _NUMERIC_TO_LETTER:
        return _NUMERIC_TO_LETTER[text]
    # Be forgiving of common near-JSON values: "B.", "(B)", "Option B".
    m = re.search(r'\b([A-D])\b', text)
    return m.group(1) if m else ''


def parse_mcq_json(text: str, qid=None) -> str:
    """
    Parse the (single) answer from a model response.

    Accepts JSON forms:
        {"id": ..., "answer": "X"}
        [{"id": ..., "answer": "X"}]
        {"id": ..., "answers": ["X"]}
        [{"id": ..., "answers": ["X"]}]

    Returns the single normalized letter ('A'/'B'/'C'/'D') for `qid` or ''
    only when both JSON and prose extraction fail.
    """
    data = parse_json_flexible(text)
    if data is not None:
        if isinstance(data, dict):
            data = [data]
        if isinstance(data, list):
            target = norm_qid(qid) if qid is not None else None
            for item in data:
                if not isinstance(item, dict):
                    continue
                this_qid = norm_qid(item.get('id', ''))
                if target is not None and this_qid != target:
                    continue
                ans = item.get('answer', item.get('answers', ''))
                letter = _normalize_letter_value(ans)
                if letter:
                    return letter
            # JSON parsed but no matching item / no letter extracted —-
            # fall through to letter-regex on the original text.
    # Prose / non-JSON fallback
    return parse_mcq_letters(text)


def parse_mcq_letters(text: str, is_multi_answer: bool = False) -> str:
    """
    Single-letter extraction from prose outputs. The `is_multi_answer` arg is
    kept for back-compat but defaults to False and is otherwise ignored —-
    the benchmark is single-choice.

    Returns 'A' / 'B' / 'C' / 'D' or '' on failure.
    """
    if not text:
        return ''
    # Strict pattern first: 'Answer: A' / 'The answer is B.'
    m = re.search(r'answer[s]?\s*[:\s]+(?:is\s+)?([A-D])', text, re.IGNORECASE)
    if m:
        return m.group(1).upper()
    # Fallback: first standalone letter A/B/C/D
    m = re.search(r'\b([A-D])\b', text.upper())
    if m:
        return m.group(1)
    # Some models reliably emit "answer: 3" instead of a letter.
    m = re.search(r'answer[s]?\s*[:\s]+(?:is\s+)?([1-4])\b', text, re.IGNORECASE)
    if m:
        return _NUMERIC_TO_LETTER[m.group(1)]
    m = re.search(r'\b([1-4])\b', text)
    return _NUMERIC_TO_LETTER[m.group(1)] if m else ''


def parse_mcq_any(text: str, qid=None) -> str:
    """Parse JSON-style responses first, then fall back to prose/letter output."""
    return parse_mcq_json(text, qid=qid) or parse_mcq_letters(text)