from __future__ import annotations
import re
from collections import Counter
from collections.abc import Mapping, Sequence
from typing import Any
import numpy as np
THINK_PROTOCOL_REVISION = 'edgerock_think_format_v1'
NATURAL_FINISH_REASONS = frozenset({'eos'})
CAP_FINISH_REASONS = frozenset({'length', 'max_tokens', 'max_new_tokens', 'max_output_tokens', 'prompt_length_cap', 'total_length_cap', 'capped'})
TIMEOUT_FINISH_REASONS = frozenset({'timeout', 'timeout_dropped', 'generation_timeout_dropped'})
THINK_IRREVERSIBLE_CAP_FAILURE_REASONS = frozenset({'redundant_think_open', 'repeated_think_close', 'empty_reasoning_before_think_close', 'think_close_text_token_mismatch'})

def think_protocol_status(response_text: str, finish_reason: str, response_ids: Sequence[int] | None=None, *, think_open_id: int | None=None, think_close_id: int | None=None, prompt_think_closed: bool=False, allow_empty_reasoning: bool=False) -> dict[str, Any]:
    text = str(response_text or '')
    finish = str(finish_reason or '').strip().lower()
    ids_available = response_ids is not None
    ids = np.asarray(response_ids if response_ids is not None else (), dtype=np.int64)
    literal_open_count = text.count('<think>')
    literal_close_count = text.count('</think>')
    token_open_count = int(np.count_nonzero(ids == int(think_open_id))) if ids_available and think_open_id is not None else None
    close_positions = np.flatnonzero(ids == int(think_close_id)).tolist() if ids_available and think_close_id is not None else None
    token_close_count = len(close_positions) if close_positions is not None else None
    opening_count = max(literal_open_count, 0 if token_open_count is None else token_open_count)
    closing_count = max(literal_close_count, 0 if token_close_count is None else token_close_count)

    def result(verdict: str, reason: str) -> dict[str, Any]:
        capped = finish in CAP_FINISH_REASONS
        irreversible_at_cap = bool(capped and verdict == 'FAIL' and (reason in THINK_IRREVERSIBLE_CAP_FAILURE_REASONS))
        return {'revision': THINK_PROTOCOL_REVISION, 'verdict': verdict, 'reason': reason, 'opening_count': opening_count, 'closing_count': closing_count, 'literal_opening_count': literal_open_count, 'literal_closing_count': literal_close_count, 'token_opening_count': token_open_count, 'token_closing_count': token_close_count, 'think_close_indices': close_positions, 'finish_reason': finish, 'finish_capped': capped, 'irreversible_at_cap': irreversible_at_cap, 'append_repairable_at_cap': bool(capped and (not irreversible_at_cap))}
    if finish in TIMEOUT_FINISH_REASONS:
        return result('EXCLUDED', 'generation_timeout_dropped')
    capped = finish in CAP_FINISH_REASONS
    natural = finish in NATURAL_FINISH_REASONS
    if not capped and (not natural):
        return result('EXCLUDED', 'generation_finish_unresolved')
    if token_close_count is not None and literal_close_count != token_close_count:
        return result('FAIL', 'think_close_text_token_mismatch')
    if opening_count:
        return result('FAIL', 'redundant_think_open')
    if closing_count > (0 if prompt_think_closed else 1):
        return result('FAIL', 'repeated_think_close')
    if closing_count == 0 and (not prompt_think_closed):
        if capped:
            return result('CAPPED_INCOMPLETE', 'generation_reached_cap')
        return result('FAIL', 'missing_think_close_at_natural_end')
    if prompt_think_closed:
        has_reasoning = True
        has_visible_answer = bool(text.strip())
    elif literal_close_count == 1:
        reasoning, visible = text.split('</think>', 1)
        visible = re.sub('[ \\t\\r\\n]*<\\|im_end\\|>[ \\t\\r\\n]*$', '', visible)
        has_reasoning = bool(reasoning.strip())
        has_visible_answer = bool(visible.strip())
    elif token_close_count == 1 and think_close_id is not None:
        close_index = int(np.flatnonzero(ids == int(think_close_id))[0])
        has_reasoning = close_index > 0
        has_visible_answer = close_index < len(ids) - 1
    else:
        return result('FAIL', 'think_close_text_token_mismatch')
    if not has_reasoning and (not allow_empty_reasoning):
        return result('FAIL', 'empty_reasoning_before_think_close')
    if not has_visible_answer:
        if capped:
            return result('CAPPED_INCOMPLETE', 'generation_reached_cap')
        return result('FAIL', 'empty_visible_answer_after_think_close')
    if capped:
        return result('CAPPED_INCOMPLETE', 'generation_reached_cap')
    return result('PASS', 'none')
PYTHON_FENCE_CONTRACT = 'visible_python_fence'
CODE_PROTOCOL_REVISION = 'visible_python_fence_append_repairable_open_prefix_v3'
CODE_IRREVERSIBLE_CAP_FAILURE_REASONS = frozenset({'empty_python_body', 'multiple_or_nested_fences', 'no_visible_python_fence', 'text_outside_python_fence'})
MECHANICAL_REPETITION_FINISH_REASONS = frozenset({'mechanical_repetition', 'repetition', 'repetition_detected'})
MECHANICAL_REPETITION_EARLY_STOP_REVISION = 'exact_tail_pattern_128to256_tokens_repeat8_first_unit_v2'
MECHANICAL_REPETITION_EARLY_STOP_MIN_PATTERN_TOKENS = 128
MECHANICAL_REPETITION_EARLY_STOP_MAX_PATTERN_TOKENS = 256
MECHANICAL_REPETITION_EARLY_STOP_MIN_REPETITIONS = 8

def first_repeated_unit_onset(response_ids: Sequence[int], end: int, maximum_period: int=256) -> int | None:
    for period in range(1, min(maximum_period, end // 2) + 1):
        if list(response_ids[end - period:end]) != list(response_ids[end - 2 * period:end - period]):
            continue
        start = end - 2 * period
        while start > 0 and response_ids[start - 1] == response_ids[start - 1 + period]:
            start -= 1
        if end - start >= max(16, 4 * period):
            return start + period
    return None

def mechanical_repetition_early_stop_patterns(tails: np.ndarray, lengths: np.ndarray, *, min_pattern_tokens=MECHANICAL_REPETITION_EARLY_STOP_MIN_PATTERN_TOKENS, max_pattern_tokens=MECHANICAL_REPETITION_EARLY_STOP_MAX_PATTERN_TOKENS, min_repetitions=MECHANICAL_REPETITION_EARLY_STOP_MIN_REPETITIONS) -> np.ndarray:
    tails = np.asarray(tails)
    lengths = np.minimum(np.asarray(lengths), tails.shape[1])
    matched = np.zeros(len(tails), dtype=np.int32)
    repeats = min_repetitions
    maximum = min(max_pattern_tokens, tails.shape[1] // repeats)
    periods = np.arange(min_pattern_tokens, maximum + 1)
    if not len(periods):
        return matched
    possible = (lengths[:, None] >= periods[None, :] * repeats) & (tails[:, -1:] == tails[:, -periods - 1])
    columns = periods > 1
    possible[:, columns] &= tails[:, -2:-1] == tails[:, -periods[columns] - 2]
    for column in np.flatnonzero(possible.any(axis=0)):
        period = int(periods[column])
        candidates = np.flatnonzero((matched == 0) & possible[:, column])
        if not len(candidates):
            continue
        blocks = tails[candidates, -period * repeats:].reshape(-1, repeats, period)
        exact = (blocks[:, :-1] == blocks[:, -1:]).all(axis=(1, 2))
        matched[candidates[exact]] = period
    return matched

def mechanical_repetition_early_stop_metrics(response_ids: Sequence[int], *, min_pattern_tokens: int=MECHANICAL_REPETITION_EARLY_STOP_MIN_PATTERN_TOKENS, max_pattern_tokens: int=MECHANICAL_REPETITION_EARLY_STOP_MAX_PATTERN_TOKENS, min_repetitions: int=MECHANICAL_REPETITION_EARLY_STOP_MIN_REPETITIONS) -> dict[str, Any]:
    if min_pattern_tokens <= 0 or max_pattern_tokens < min_pattern_tokens:
        raise ValueError('mechanical repetition pattern range is invalid')
    if min_repetitions < 2:
        raise ValueError('mechanical repetition requires at least two repeats')
    response_tokens = len(response_ids)
    values = np.asarray(response_ids[-max_pattern_tokens * min_repetitions:], dtype=np.int64)
    pattern_tokens = int(mechanical_repetition_early_stop_patterns(values[None, :], np.array([len(values)]), min_pattern_tokens=min_pattern_tokens, max_pattern_tokens=max_pattern_tokens, min_repetitions=min_repetitions)[0])
    if pattern_tokens:
        repeated_tokens = pattern_tokens * min_repetitions
        return {'revision': MECHANICAL_REPETITION_EARLY_STOP_REVISION, 'stop': True, 'response_tokens': response_tokens, 'pattern_tokens': pattern_tokens, 'minimum_repetitions': min_repetitions, 'repeated_tokens': repeated_tokens, 'repeated_tail_start': response_tokens - repeated_tokens, 'failure_onset_index': first_repeated_unit_onset(response_ids, response_tokens, pattern_tokens), 'confirmation_index': response_tokens - 1}
    return {'revision': MECHANICAL_REPETITION_EARLY_STOP_REVISION, 'stop': False, 'response_tokens': response_tokens, 'pattern_tokens': None, 'minimum_repetitions': min_repetitions, 'repeated_tokens': 0, 'repeated_tail_start': None, 'failure_onset_index': None, 'confirmation_index': None}

def mechanical_repetitive_tail_metrics(response_ids: Sequence[int], finish_reason: str, *, tail_tokens: int=256, require_exact_confirmation: bool=False) -> dict[str, Any]:
    values = response_ids
    finish = str(finish_reason).strip().lower()
    capped = finish in CAP_FINISH_REASONS
    generation_stopped_repetition = finish in MECHANICAL_REPETITION_FINISH_REASONS
    if tail_tokens < 16:
        raise ValueError('degenerate-tail window must contain at least 16 tokens')
    tail = list(map(int, values[-tail_tokens:]))
    fourgrams = [tuple(tail[index:index + 4]) for index in range(len(tail) - 3)]
    unique_token_ratio = len(set(tail)) / len(tail) if tail else 1.0
    repeated_fourgram_ratio = 1.0 - len(set(fourgrams)) / len(fourgrams) if fourgrams else 0.0
    eligible_tail = len(tail) >= 16
    exact = None
    if not generation_stopped_repetition and (require_exact_confirmation or not (eligible_tail and repeated_fourgram_ratio >= 0.9)):
        for end in dict.fromkeys((len(values), len(values) - 1, len(values) // 128 * 128)):
            if end >= 1024:
                candidate = mechanical_repetition_early_stop_metrics(values if end == len(values) else values[:end])
                if candidate['stop']:
                    exact = candidate
                    break
    repetitive_degenerate = bool(generation_stopped_repetition or exact is not None or (not require_exact_confirmation and eligible_tail and (repeated_fourgram_ratio >= 0.9)))
    onset: int | None = None
    confirmation: int | None = None
    if generation_stopped_repetition:
        confirmation = len(values) - 1 if len(values) else None
        onset = first_repeated_unit_onset(values, len(values))
    elif exact is not None:
        onset, confirmation = (exact['failure_onset_index'], exact['confirmation_index'])
    elif repetitive_degenerate:
        values = list(map(int, values))
        first_end = 15
        grams = Counter((tuple(values[index:index + 4]) for index in range(0, first_end - 2)))
        for end in range(first_end, len(values)):
            if end > first_end:
                added_start = end - 3
                grams[tuple(values[added_start:added_start + 4])] += 1
                if end + 1 > tail_tokens:
                    removed_start = end - tail_tokens
                    removed = tuple(values[removed_start:removed_start + 4])
                    grams[removed] -= 1
                    if grams[removed] == 0:
                        del grams[removed]
            width = min(end + 1, tail_tokens) - 3
            if 1.0 - len(grams) / width >= 0.9:
                confirmation = end
                onset = first_repeated_unit_onset(values, end + 1)
                break
        if confirmation is None:
            raise RuntimeError('mechanical failure lacks a threshold-crossing token')
    return {'capped': capped, 'response_tokens': len(values), 'tail_tokens': len(tail), 'unique_token_ratio': unique_token_ratio, 'repeated_fourgram_ratio': repeated_fourgram_ratio, 'generation_stopped_repetition': generation_stopped_repetition, 'repetitive_degenerate': repetitive_degenerate, 'degenerate': repetitive_degenerate, 'failure_onset_index': onset, 'confirmation_index': confirmation}

def _contract_name(output_contract: str | Mapping[str, Any]) -> str:
    if isinstance(output_contract, Mapping):
        return str(output_contract.get('output_contract') or '').strip().lower()
    return str(output_contract or '').strip().lower()

def _opening_fence_prefix_is_append_repairable(visible: str) -> bool:
    prefix = str(visible).lstrip()
    if not prefix:
        return True
    if '\n' in prefix:
        return False
    pending_crlf = prefix.endswith('\r') and '\r' not in prefix[:-1]
    if pending_crlf:
        prefix = prefix[:-1]
    elif '\r' in prefix:
        return False
    if len(prefix) < 3:
        return prefix == '```'[:len(prefix)]
    if not prefix.startswith('```'):
        return False
    language = prefix[3:]
    stripped = language.strip().lower()
    if pending_crlf or language != language.rstrip():
        return stripped in {'', 'py', 'python'}
    return any((value.startswith(stripped) for value in ('', 'py', 'python')))

def code_contract_status(response_text: str, response_ids: Sequence[int] | None, finish_reason: str, output_contract: str | Mapping[str, Any], *, eos_token_ids: set[int] | None=None, prompt_think_closed: bool=False) -> dict[str, Any]:
    contract = _contract_name(output_contract)
    finish = str(finish_reason or '').strip().lower()
    finish_natural = finish in NATURAL_FINISH_REASONS
    finish_capped = finish in CAP_FINISH_REASONS
    if contract != PYTHON_FENCE_CONTRACT:
        reason = 'output_contract_missing' if not contract else 'unsupported_output_contract'
        return {'revision': CODE_PROTOCOL_REVISION, 'applicable': True, 'contract_pass': False, 'contract_verdict': 'FAIL', 'contract_reason': reason, 'structure_reason': reason, 'visible_final': '', 'visible_code': '', 'opening_language': '', 'finish_natural': finish_natural, 'single_unterminated_python_fence': False, 'eos_after_contract': None}
    text = str(response_text or '')
    status: dict[str, Any] = {'revision': CODE_PROTOCOL_REVISION, 'applicable': True, 'contract_pass': False, 'contract_verdict': 'FAIL', 'contract_reason': '', 'structure_reason': '', 'visible_final': '', 'visible_code': '', 'opening_language': '', 'finish_natural': finish_natural, 'finish_capped': finish_capped, 'irreversible_at_cap': False, 'append_repairable_at_cap': finish_capped, 'single_unterminated_python_fence': False, 'eos_after_contract': None}

    def fail(reason: str, *, structure_reason: str | None=None) -> dict[str, Any]:
        structural = structure_reason or reason
        status['contract_reason'] = reason
        status['structure_reason'] = structural
        status['irreversible_at_cap'] = bool(finish_capped and structural in CODE_IRREVERSIBLE_CAP_FAILURE_REASONS)
        status['append_repairable_at_cap'] = bool(finish_capped and (not status['irreversible_at_cap']))
        if status['append_repairable_at_cap']:
            status['contract_verdict'] = 'CAPPED_INCOMPLETE'
        return status
    if '</think>' not in text and (not prompt_think_closed):
        if finish_capped:
            return fail('generation_open', structure_reason='generation_open')
        return fail('no_visible_python_fence')
    visible = text if prompt_think_closed else text.split('</think>', 1)[1]
    visible = re.sub('[ \\t\\r\\n]*<\\|im_end\\|>[ \\t\\r\\n]*$', '', visible).lstrip()
    status['visible_final'] = visible.strip()
    if any((marker in visible for marker in ('<|im_start|>', '<|assistant|>', '<|user|>', '<|system|>'))):
        return fail('text_outside_python_fence')
    opening = re.match('^```([^\\n`]*)\\r?\\n', visible)
    if opening is None or opening.group(1).strip().lower() not in {'', 'py', 'python'}:
        structure_reason = 'no_visible_python_fence'
        if finish_capped and _opening_fence_prefix_is_append_repairable(visible):
            return fail('generation_open', structure_reason='unterminated_python_fence')
        reason = 'generation_open' if not finish_natural else structure_reason
        return fail(reason, structure_reason=structure_reason)
    status['opening_language'] = opening.group(1).strip().lower()
    fences = [match.start() for match in re.finditer('```', visible)]
    if len(fences) == 1:
        status['visible_code'] = visible[opening.end():].strip()
        status['single_unterminated_python_fence'] = True
        structure_reason = 'unterminated_python_fence'
        reason = 'generation_open' if not finish_natural else structure_reason
        return fail(reason, structure_reason=structure_reason)
    if len(fences) != 2:
        return fail('multiple_or_nested_fences')
    closing_start = fences[1]
    code = visible[opening.end():closing_start]
    trailing = visible[closing_start + 3:]
    status['visible_code'] = code.strip()
    if trailing.strip():
        return fail('text_outside_python_fence')
    if not code.strip():
        return fail('empty_python_body')
    if not finish_natural:
        return fail('generation_open', structure_reason='none')
    if response_ids is not None and eos_token_ids is not None:
        ids = list(map(int, response_ids))
        eos_positions = [index for index, token in enumerate(ids) if token in eos_token_ids]
        status['eos_after_contract'] = bool(eos_positions and eos_positions[-1] == len(ids) - 1)
        if not status['eos_after_contract']:
            return fail('eos_before_contract_complete', structure_reason='none')
    status.update({'contract_pass': True, 'contract_verdict': 'PASS', 'contract_reason': 'none', 'structure_reason': 'none'})
    return status