from __future__ import annotations

import re

from albedo_eval_service.shared.observation_format import first_bash_block

_THOUGHT_LINE_RE = re.compile(r"^\s*THOUGHT:\s*(.*)$", re.IGNORECASE)
_THINK_BLOCK_RE = re.compile(r"<think>(.*?)</think>", re.DOTALL | re.IGNORECASE)
_FENCE_RE = re.compile(r"```[\w-]*[ \t]*\n.*?```", re.DOTALL)


def thought_text(gold: str) -> str:
    """Pull a thinking body out of rendered gold (`THOUGHT:` or `<think>`)."""
    text = gold or ""
    blocks = _THINK_BLOCK_RE.findall(text)
    if blocks:
        return blocks[0].strip()
    lines: list[str] = []
    taking = False
    for line in text.splitlines():
        match = _THOUGHT_LINE_RE.match(line)
        if match:
            taking = True
            if match.group(1).strip():
                lines.append(match.group(1).strip())
            continue
        if taking:
            if line.strip().startswith("```") or line.strip().startswith("<"):
                break
            lines.append(line)
    if lines:
        return "\n".join(lines).strip()
    before = _FENCE_RE.split(text, maxsplit=1)[0]
    before = _THINK_BLOCK_RE.sub("", before)
    cleaned = re.sub(r"^\s*THOUGHT:\s*", "", before, flags=re.I).strip()
    return cleaned


def wrap_completion(gold: str, bash_override: str | None = None) -> str | None:
    """One closed `<think>` block plus exactly one bash fence. Live thinking is on."""
    bash = (bash_override or first_bash_block(gold) or "").strip()
    if not bash:
        return None
    thought = thought_text(gold)
    if not thought:
        thought = _fallback_thought(bash)
    thought = thought.replace("</think>", "").strip()
    return f"<think>\n{thought}\n</think>\n\n```bash\n{bash}\n```\n"


def _fallback_thought(bash: str) -> str:
    head = bash.splitlines()[0][:160]
    if bash.startswith("echo ") and "SUBMIT" in bash.upper():
        return f"Work is saved. Submit with the exact command from the instructions: {head}"
    return f"Next step is to run this command and inspect the result: {head}"