from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import Any, Dict, List, Tuple

import numpy as np


_INT_RE = re.compile(r"-?\d+")
_TUPLE_PROMPT_RE = re.compile(r"\((\d+),(\d+),(\d+)\)")


@dataclass(frozen=True)
class CellExample:
    grid: np.ndarray
    target_cell: tuple[int, int]
    target_value: int
    turn_idx: int
    total_turns: int


def parse_n_value_prediction(text: str, n: int) -> Tuple[List[int] | None, bool]:
    raw = str(text or '').strip()
    if not raw:
        return None, False

    try:
        parsed = json.loads(raw)
        if isinstance(parsed, dict) and isinstance(parsed.get('values'), list):
            values = [int(v) for v in parsed['values']]
            if len(values) == int(n):
                return values, True
        if isinstance(parsed, list):
            values = [int(v) for v in parsed]
            if len(values) == int(n):
                return values, True
    except Exception:
        pass

    values = [int(match.group(0)) for match in _INT_RE.finditer(raw)]
    if len(values) == int(n):
        return values, True
    return None, False


def parse_grid_from_tuple_prompt(prompt_text: str) -> np.ndarray:
    triples = _TUPLE_PROMPT_RE.findall(str(prompt_text))
    if len(triples) < 16:
        raise ValueError('Could not recover 16 (row,col,value) tuples from prompt.')
    grid = np.zeros((4, 4), dtype=int)
    for rr, cc, vv in triples[:16]:
        grid[int(rr) - 1, int(cc) - 1] = int(vv)
    return grid


def build_cell_examples_from_row(row: Dict[str, Any]) -> List[CellExample]:
    prompt = str(row['prompt'])
    grid = parse_grid_from_tuple_prompt(prompt)
    metadata = dict(row.get('metadata', {}))
    empty_locs = metadata.get('empty_locs_1based')
    target_triples = metadata.get('target_triples_1based')

    if not empty_locs or not target_triples:
        completion = str(row.get('completion', ''))
        parsed, _ = parse_n_value_prediction(completion, int(metadata.get('empties', 0) or 0))
        if parsed is None:
            raise ValueError('Row is missing metadata and completion could not be parsed.')
        empty_locs = [(r + 1, c + 1) for r, c in np.argwhere(grid == 0).tolist()]
        target_triples = [(int(r), int(c), int(v)) for (r, c), v in zip(empty_locs, parsed)]

    total_turns = len(target_triples)
    out: List[CellExample] = []
    for idx, triple in enumerate(target_triples, start=1):
        rr, cc, value = int(triple[0]) - 1, int(triple[1]) - 1, int(triple[2])
        out.append(
            CellExample(
                grid=np.asarray(grid, dtype=int).copy(),
                target_cell=(rr, cc),
                target_value=value,
                turn_idx=idx,
                total_turns=total_turns,
            )
        )
    return out