import asyncio
from typing import Any, Sequence

from tinker_cookbook.completers import TokenCompleter
from tinker_cookbook.rl.types import (
    Env,
    EnvGroupBuilder,
    Trajectory,
    TrajectoryGroup,
    Transition,
)
from tinker_cookbook.utils import logtree

# Max characters for log values in table cells before truncation
LOG_VALUE_MAX_LEN = 100


def _truncate_log_value(value: Any, max_len: int = LOG_VALUE_MAX_LEN) -> tuple[str, bool]:
    """Truncate a log value if it's too long. Returns (display_value, was_truncated)."""
    str_value = str(value)
    if len(str_value) > max_len:
        return str_value[:max_len] + "...", True
    return str_value, False


# Global counter for debug prints (only print first few)
_DEBUG_ROLLOUT_COUNT = 0
_DEBUG_ROLLOUT_MAX = 3


@logtree.scope_header_decorator
async def do_single_rollout(policy: TokenCompleter, env: Env) -> Trajectory:
    global _DEBUG_ROLLOUT_COUNT
    transitions = []
    ob, stop_condition = await env.initial_observation()

    # Debug: print observation tokens to verify prefill
    if _DEBUG_ROLLOUT_COUNT < _DEBUG_ROLLOUT_MAX:
        _DEBUG_ROLLOUT_COUNT += 1
        all_tokens = []
        for chunk in ob.chunks:
            if hasattr(chunk, "tokens"):
                all_tokens.extend(chunk.tokens)
        print(f"\n{'=' * 60}")
        print(f"[DEBUG ROLLOUT #{_DEBUG_ROLLOUT_COUNT}] Observation tokens (last 30):")
        print(f"  Tokens: {all_tokens[-30:]}")
        print(f"  Total prompt length: {len(all_tokens)} tokens")
        print(f"{'=' * 60}\n")

    while True:
        ac_with_logprobs = await policy(ob, stop_condition)
        step_result = await env.step(ac_with_logprobs.tokens)
        transition = Transition(
            ob=ob,
            ac=ac_with_logprobs,
            reward=step_result.reward,
            episode_done=step_result.episode_done,
            metrics=step_result.metrics,
            logs=step_result.logs,
        )
        transitions.append(transition)
        ob = step_result.next_observation
        stop_condition = step_result.next_stop_condition
        if step_result.episode_done:
            break
    return Trajectory(transitions=transitions, final_ob=ob)


@logtree.scope_header_decorator
async def do_group_rollout(
    env_group_builder: EnvGroupBuilder, policy: TokenCompleter
) -> TrajectoryGroup:
    envs_G: Sequence[Env] = await env_group_builder.make_envs()
    trajectories_G = await asyncio.gather(*[do_single_rollout(policy, env) for env in envs_G])
    rewards_and_metrics_G = await env_group_builder.compute_group_rewards(trajectories_G, envs_G)
    rewards_G, metrics_G = zip(*rewards_and_metrics_G, strict=True)

    # Log trajectory tables with final rewards
    with logtree.scope_header("Trajectory Summary"):
        for i, (traj, final_reward) in enumerate(zip(trajectories_G, rewards_G, strict=True)):
            # Pre-scan to collect all log keys across all transitions (preserving order, deduped)
            all_log_keys = list(dict.fromkeys(key for t in traj.transitions for key in t.logs))

            rows = []
            truncated_values: list[tuple[int, str, str]] = []  # (step, key, full_value)
            step_reward_sum = 0.0
            for t_idx, t in enumerate(traj.transitions):
                step_reward_sum += t.reward
                row: dict[str, Any] = {
                    "step": t_idx,
                    "ob_len": t.ob.length,
                    "ac_len": len(t.ac.tokens),
                    "reward": f"{t.reward:.3f}",
                }
                # Add log fields (user is responsible for avoiding collision with core columns)
                for key in all_log_keys:
                    if key in t.logs:
                        display_val, was_truncated = _truncate_log_value(t.logs[key])
                        row[key] = display_val
                        if was_truncated:
                            truncated_values.append((t_idx, key, str(t.logs[key])))
                    else:
                        row[key] = "-"
                rows.append(row)
            # Add final row with final observation and computed reward
            rows.append(
                {
                    "step": "final",
                    "ob_len": traj.final_ob.length,
                    "ac_len": "-",
                    "reward": f"{final_reward:.3f}",
                    **{key: "-" for key in all_log_keys},
                }
            )
            # Add total reward row
            rows.append(
                {
                    "step": "total",
                    "ob_len": "-",
                    "ac_len": "-",
                    "reward": f"{step_reward_sum + final_reward:.3f}",
                    **{key: "-" for key in all_log_keys},
                }
            )
            logtree.table(rows, caption=f"Trajectory {i}")

            # Show full content for any truncated values in collapsible blocks
            for step_idx, key, full_value in truncated_values:
                logtree.details(
                    full_value,
                    summary=f"Step {step_idx} - {key} (full, {len(full_value)} chars)",
                    pre=True,
                )

    return TrajectoryGroup(trajectories_G, list(rewards_G), list(metrics_G))
