"""PUP-T for one prompt: render, read the text back, rewrite the caption and render again (no training).

Same decisions as the measured research loop (p5/klein_rl/text_loop.py, wave 4):
- round 0 renders the PUP caption with the user seed;
- every candidate is read by the verifier and scored deterministically against the quoted texts of the prompt;
- the best candidate so far ranks by (1-CER, exact match), and ties keep the earliest;
- a prompt stops when its best candidate is solved (exact match and 1-CER = 1), when the rewrite is unusable, or
  when the last round did not improve on the candidate already rewritten;
- otherwise the best caption is rewritten from the expected vs observed text, and CANDIDATES_PER_ROUND candidates
  are rendered with seed, seed+1, ..., at most MAX_REFINES times.
Prompts without quoted text get the round-0 render only.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field

from .quotes import enforce_quote_guardrail, extract_quoted_literals
from .typography import parse_vlm_ocr_readout, score_typography_ocr

MAX_REFINES = 2
CANDIDATES_PER_ROUND = 2
CAPTION_MIN_CHARS = 20
CAPTION_MAX_CHARS = 4000
EPS = 1e-9

REVISE_TASK = """An image generator was given this caption:
{caption}

The image must show exactly these texts: {expected}
A careful reading of the generated image found: {observed}

Rewrite the caption so the generator renders every required text exactly, letter for letter: keep the scene and
everything else that is requested; put each required text in double quotes with its exact spelling, casing and
punctuation; place it prominently on a large plain high-contrast surface in a simple bold sans-serif font; state that each
required text appears exactly once and nowhere else, with no other text in the image; keep every other detail of the caption and do not shorten it.
Answer with one JSON object only: {{"revised_caption": "..."}}"""


@dataclass
class Candidate:
    round: int
    cand: int
    caption: str
    seed: int
    image: object
    observed: list
    ocr_reward_01: float
    exact_match_rate: float


@dataclass
class Result:
    image: object
    best: Candidate | None
    caption: str
    expected: list
    candidates: list = field(default_factory=list)
    revisions: list = field(default_factory=list)
    stop: str = ''

    @property
    def renders(self) -> int:
        return max(1, len(self.candidates))


def parse_revision(text: str) -> str:
    """revised_caption from the verifier's JSON; '' when absent, malformed or out of bounds."""
    decoder = json.JSONDecoder()
    text = text or ''
    pos = text.find('{')
    while pos >= 0:
        try:
            obj, _end = decoder.raw_decode(text, pos)
        except (ValueError, RecursionError):
            obj = None
        if isinstance(obj, dict) and 'revised_caption' in obj:
            value = obj['revised_caption']
            value = value.strip() if isinstance(value, str) else ''
            return value if CAPTION_MIN_CHARS <= len(value) <= CAPTION_MAX_CHARS else ''
        pos = text.find('{', pos + 1)
    return ''


def _rank(c: Candidate):
    return (c.ocr_reward_01, c.exact_match_rate >= 1.0 - EPS)


def _solved(c: Candidate) -> bool:
    return c.exact_match_rate >= 1.0 - EPS and c.ocr_reward_01 >= 1.0 - EPS


def _best(candidates):
    best = None
    for c in candidates:  # rounds and candidates in order, so ties keep the earliest
        if best is None or _rank(c) > _rank(best):
            best = c
    return best


def caption_for(prompt: str, upsample) -> str:
    """PUP caption with every quoted literal of the prompt kept verbatim (the measured guardrail)."""
    repaired, _check = enforce_quote_guardrail(prompt, upsample(prompt))
    return repaired


def generate(prompt: str, seed: int, *, upsample, render, read=None, revise=None, caption: str | None = None,
             max_refines: int = MAX_REFINES) -> Result:
    """upsample(prompt) -> caption; render(caption, seed) -> image; read(image) -> raw OCR reply;
    revise(caption, expected, observed) -> raw JSON reply. read/revise are needed only for prompts with quoted text."""
    expected = list(extract_quoted_literals(prompt))
    if int(max_refines) < 0:
        raise ValueError('max_refines must be >= 0')
    if caption is None and upsample is None:
        raise ValueError('upsample is required unless a caption is given')
    if expected and (read is None or revise is None):
        raise ValueError('prompts with quoted text need read and revise')
    caption = caption if caption is not None else caption_for(prompt, upsample)
    image0 = render(caption, seed)
    if not expected:
        return Result(image=image0, best=None, caption=caption, expected=[], stop='no_quoted_text')

    def check(rnd, cand, cap, s, image):
        observed = parse_vlm_ocr_readout(read(image))
        score = score_typography_ocr(expected, observed)
        return Candidate(rnd, cand, cap, s, image, observed, float(score['ocr_reward_01']),
                         float(score['exact_match_rate']))

    result = Result(image=image0, best=None, caption=caption, expected=expected)
    result.candidates.append(check(0, 0, caption, seed, image0))
    for rnd in range(max_refines):
        best = _best(result.candidates)
        if _solved(best):
            result.stop = 'solved'
            break
        if best.round != rnd:
            result.stop = 'no_improvement'
            break
        revised = parse_revision(revise(best.caption, expected, best.observed))
        result.revisions.append({'round': rnd + 1, 'from': [best.round, best.cand], 'revised_caption': revised})
        if not revised:
            result.stop = 'unusable_rewrite'
            break
        for j in range(CANDIDATES_PER_ROUND):
            result.candidates.append(check(rnd + 1, j, revised, seed + j, render(revised, seed + j)))
    else:
        result.stop = 'solved' if _solved(_best(result.candidates)) else 'max_refines'
    result.best = _best(result.candidates)
    result.image = result.best.image
    return result