"""带「反事实编辑信用」的 fully-async 训练器:在 SkyRL 之外(recipe 内)做子类,不改 rl/skyrl/SkyRL。

只做两件事(评审 §B7):
  1. convert_to_training_input:父类照常建批,然后把 generator 给的每行 `credit_advantages`(response 局部坐标、已乘 λ·D·s)
     右对齐贴成张量(response 张量在 preprocess 里就是右对齐的),乘 loss_mask;
  2. compute_advantages_and_returns:父类照常算 GRPO 优势(组均值基线、广播),然后 advantages += credit(out-of-place,returns 不动),
     再由父类后续的 prompt-mean 缩放 / DPPO 掩码照常处理。

credit 全 0 时与父类逐位等价(臂 O 和臂 A 重启都走这条路)。指标前缀 credit/。
两个核心步骤抽成模块级函数(attach_credit / add_credit),便于不起训练器就能单测。
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Sequence
import json
from pathlib import Path

import torch
from loguru import logger
from swe_skyrl_recipe.credit_readiness import ConsumedBatchHealth
from swe_skyrl_recipe.terminal_edit_focus import build_terminal_edit_focus

CREDIT_KEY = "credit_advantages"


def attach_credit(loss_mask: torch.Tensor, response_mask: torch.Tensor, rows: Optional[Sequence[Optional[Sequence[float]]]]) -> torch.Tensor:
    """rows[i] = 第 i 条 response 的信用(长度 = 该条 response token 数)。返回与 loss_mask 同形的 float32 张量,右对齐、乘 loss_mask。"""
    credit = torch.zeros(loss_mask.shape, dtype=torch.float32, device=loss_mask.device)
    if rows is None:
        return credit
    width = credit.shape[1]
    for i in range(min(len(rows), credit.shape[0])):
        row = rows[i]
        n = len(row) if row is not None else 0
        if n == 0:
            continue
        if n > width:
            raise ValueError(f"credit 行 {i} 长 {n} > response 宽 {width}:generator 截断没同步")
        # 严格:credit 行长必须等于这条 response 的有效 token 数(短行/长行都拒绝;0910 独立审计:短行会被静默右移)
        n_resp = int(response_mask[i].sum().item())
        if n != n_resp:
            raise ValueError(f"credit 行 {i}:长度 {n} != response 有效 token 数 {n_resp}")
        credit[i, width - n :] = torch.as_tensor(list(row), dtype=torch.float32, device=credit.device)
    return credit * loss_mask.to(credit.dtype)


def add_credit(advantages: torch.Tensor, credit: torch.Tensor) -> tuple[torch.Tensor, Dict[str, float]]:
    """advantages + credit(out-of-place),并给出 credit/ 指标。"""
    if credit.shape != advantages.shape:
        raise ValueError(f"credit 形状 {tuple(credit.shape)} != advantages {tuple(advantages.shape)}")
    credit = credit.to(advantages.device, advantages.dtype)
    combined = advantages + credit
    nz = credit != 0
    n_tok = int(nz.sum().item())
    metrics: Dict[str, float] = {
        "credit/tokens_with_credit": float(n_tok),
        "credit/rows_with_credit": float(nz.any(dim=1).sum().item()),
        "credit/signed_mass": float(credit.sum().item()),
        "credit/abs_mass": float(credit.abs().sum().item()),
        "credit/max_abs": float(credit.abs().max().item()) if n_tok else 0.0,
    }
    if n_tok:
        a, c = advantages[nz], combined[nz]
        metrics["credit/sign_flip_frac"] = float(((torch.sign(a) != torch.sign(c)) & (a != 0)).float().mean().item())
        metrics["credit/abs_ratio_to_terminal"] = float((credit[nz].abs().mean() / (a.abs().mean() + 1e-8)).item())
    return combined, metrics


try:
    from skyrl.train.fully_async_trainer import FullyAsyncRayPPOTrainer
except Exception:  # noqa: BLE001 —— 单测环境没有 SkyRL 时只用上面两个函数
    FullyAsyncRayPPOTrainer = object  # type: ignore[misc,assignment]


class CreditFullyAsyncRayPPOTrainer(FullyAsyncRayPPOTrainer):  # type: ignore[misc]
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        alg = self.cfg.trainer.algorithm
        assert alg.advantage_estimator == "grpo", f"credit 训练器只在 grpo 优势下验证过,当前 {alg.advantage_estimator}"
        if getattr(alg, "advantage_batch_normalize", False):
            raise ValueError("advantage_batch_normalize 会把辅助信用一起白化,评审规格禁止;请关掉")
        self._credit_mode = str(self.cfg.generator.edit_credit_mode or 'off').lower()
        self._health_enabled = bool(getattr(self.cfg.generator, 'edit_credit_health_checks', False))
        self._batch_health = ConsumedBatchHealth(self._credit_mode)
        self._credit_ledger = Path(self.cfg.trainer.ckpt_path).parent / 'consumed-credit.jsonl'
        logger.info("CreditFullyAsyncRayPPOTrainer: 已启用(credit 全 0 时与 FullyAsyncRayPPOTrainer 逐位等价)")

    def convert_to_training_input(self, generator_output, uids):
        rows = generator_output.get(CREDIT_KEY) if isinstance(generator_output, dict) else None
        data = super().convert_to_training_input(generator_output, uids)
        data[CREDIT_KEY] = attach_credit(data["loss_mask"], data["response_mask"], rows)
        if self._credit_mode == 'terminal_edit':
            data['terminal_edit_mask'] = attach_credit(data['loss_mask'], data['response_mask'], generator_output['terminal_edit_mask'])
        if self._health_enabled:
            records = generator_output.get('credit_observations')
            if records is None or len(records) != len(uids):
                raise ValueError('Missing per-row credit observations for successor run')
            data.metadata['credit_observations'] = records
        return data

    def _write_credit_event(self, phase, **details):
        if self._health_enabled:
            self._credit_ledger.parent.mkdir(parents=True, exist_ok=True)
            with self._credit_ledger.open('a') as f:
                f.write(json.dumps(dict(phase=phase, global_step=self.global_step, mode=self._credit_mode, **details)) + '\n')

    async def _run_training(self, training_input):
        if self._health_enabled:
            records = training_input.metadata['credit_observations']
            live = training_input['loss_mask'][:len(records)].bool().any(dim=1).tolist()
            health = self._batch_health.observe(records, live)
            self._write_credit_event('selected_for_update', health=health, rows=records)
            for key in ('rows', 'live_rows', 'cf_issued', 'cf_usable'):
                self.all_metrics['credit/consumed_' + key] = health[key]
            if health['stop_reasons']:
                self._write_credit_event('stopped_before_update', reasons=health['stop_reasons'])
                raise RuntimeError('Credit batch health stop: ' + ', '.join(health['stop_reasons']))
        status = await super()._run_training(training_input)
        # The backend does not expose every optimizer's skipped-step boolean.
        # A normal return proves completion of this call, not weight mutation.
        self._write_credit_event('training_call_completed', optimizer_update_success_verified=False,
                                 returned_status={k:float(v) for k,v in status.items() if isinstance(v, (int,float))})
        return status

    def compute_advantages_and_returns(self, data):
        credit = data.pop(CREDIT_KEY) if CREDIT_KEY in data else None
        selected = data.pop('terminal_edit_mask') if 'terminal_edit_mask' in data else None
        data = super().compute_advantages_and_returns(data)
        if self._credit_mode == 'terminal_edit':
            if selected is None:
                raise ValueError('E requires its preselected generated-token mask')
            credit = build_terminal_edit_focus(data['advantages'], data['loss_mask'], selected,
                float(self.cfg.generator.edit_credit_lambda), float(self.cfg.generator.edit_credit_token_cap))
        if credit is None:
            return data
        combined, metrics = add_credit(data["advantages"], credit)
        data["advantages"] = combined  # out-of-place:上游 GRPO 里 advantages 与 returns 是同一个张量的别名
        self.all_metrics.update(metrics)
        self._write_credit_event('advantages_prepared', metrics=metrics)
        return data