from __future__ import annotations

import os
import sys
from typing import Any, Optional

import numpy as np

CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
PARENT_DIR = os.path.dirname(CURRENT_DIR)
if PARENT_DIR not in sys.path:
    sys.path.insert(0, PARENT_DIR)

from format_utils_icon import grid_to_text


def make_multi_output_system_prompt(*, stage_i: int, total_empties_hint: int = 10) -> str:
    i = max(1, int(stage_i))
    return (
        "You are a Sudoku value policy.\n"
        f"This setup uses puzzles with about {int(total_empties_hint)} empty cells.\n"
        "You will be given one target empty cell.\n"
        'Return ONLY one JSON object of the form {"values":[...]}.\n'
        'The JSON object must contain exactly one key named "values".\n'
        'The "values" field must be a JSON array of unique integers in [1,9].\n'
        "You may return as many candidate values as you want, including one, several, or many values.\n"
        "Choose the number of returned values yourself based on which values seem i-consistent.\n"
        "The order of the values does not matter.\n"
        "Do not output any explanation, markdown, punctuation outside JSON, or extra text.\n"
        f"Current stage objective: i={i} consistency.\n"
    )


def build_multi_output_cell_prompt(
    grid_9x9: np.ndarray,
    *,
    target_cell: tuple[int, int],
    stage_i: int,
    tokenizer: Any,
    turn_idx: int,
    total_turns: int,
    prev_output_flag: Optional[str] = None,
    total_empties_hint: int = 10,
) -> str:
    g = np.asarray(grid_9x9, dtype=int).reshape(9, 9)
    empties = int(np.sum(g == 0))
    rr, cc = int(target_cell[0]), int(target_cell[1])
    system_msg = make_multi_output_system_prompt(
        stage_i=stage_i, total_empties_hint=total_empties_hint
    ).strip()
    empty_locs = [(int(r) + 1, int(c) + 1) for r, c in np.argwhere(g == 0).tolist()]
    empty_locs_text = ", ".join(f"({r},{c})" for r, c in empty_locs)
    user_msg = (
        "Sudoku grid (0 means empty):\n"
        + grid_to_text(g)
        + "\n"
        + f"Empty cells in row-major order ({empties} total): {empty_locs_text}\n\n"
        + f"Target cell to fill now: ({rr + 1},{cc + 1}).\n"
        + f"Turn: {int(turn_idx)}/{int(total_turns)}.\n"
        + 'Return only JSON with candidate values for this target cell: {"values":[...]}'
    )
    if prev_output_flag is not None:
        user_msg += f"\nPrevious output_flag (context only): {prev_output_flag}"

    chat_template = getattr(tokenizer, "chat_template", None)
    if chat_template:
        messages = [
            {"role": "system", "content": system_msg},
            {"role": "user", "content": user_msg},
        ]
        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    return system_msg + "\n\n" + user_msg + "\n"