"""
Shared preprocessing utilities for detector-to-RL pipeline.

Used by both train_intervention.py and evaluate.py to avoid circular imports.
"""

import numpy as np
import torch
from typing import List, Dict
from transformers import PreTrainedTokenizer

from src.models.detector import CompanionRiskDetector
from src.data.dataset import format_conversation, validate_and_normalize
from src.utils.taxonomy import (
    ACTION_NAME_TO_ID,
    NUM_RISK_LEVELS,
    NUM_PRIMARY,
    PRIMARY_CATEGORY_LIST,
)


def encode_sample(
    sample: Dict,
    tokenizer: PreTrainedTokenizer,
    max_persona_len: int = 128,
    max_context_len: int = 512,
    max_response_len: int = 256,
    max_history_turns: int = 5,
    device: str = "cpu",
):
    """Tokenize a single sample into three encoder inputs."""
    texts = format_conversation(
        sample["persona"],
        sample["history"],
        sample["user_input"],
        sample["ai_response"],
        max_history_turns=max_history_turns,
    )

    def enc(text: str, max_len: int) -> Dict[str, torch.Tensor]:
        return tokenizer(
            text,
            max_length=max_len,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        )

    p_enc = enc(texts["persona_text"], max_persona_len)
    c_enc = enc(texts["context_text"], max_context_len)
    r_enc = enc(texts["response_text"], max_response_len)

    return (
        p_enc["input_ids"].to(device),
        p_enc["attention_mask"].to(device),
        c_enc["input_ids"].to(device),
        c_enc["attention_mask"].to(device),
        r_enc["input_ids"].to(device),
        r_enc["attention_mask"].to(device),
    )


def preprocess_samples_with_detector(
    samples: List[Dict],
    detector: CompanionRiskDetector,
    tokenizer: PreTrainedTokenizer,
    device: str = "cpu",
    max_persona_len: int = 128,
    max_context_len: int = 512,
    max_response_len: int = 256,
    max_history_turns: int = 5,
    binary_threshold: float = 0.5,
) -> List[Dict]:
    """
    Run the detector on all samples and attach detector outputs as RL state fields.

    Adds to each sample:
      d_score        : float, risk probability from detector
      l_risk         : int, predicted risk level (overrides label if already present)
      c_primary_probs: List[float] of length NUM_PRIMARY
      c_primary_idx  : int, predicted primary category index
      e_H_pool       : List[float] of length hidden_size — context embedding
      e_P_pool       : List[float] of length hidden_size — persona embedding
    """
    detector.eval()
    processed = []

    for i, raw_sample in enumerate(samples):
        sample = validate_and_normalize(dict(raw_sample))

        ids = encode_sample(
            sample, tokenizer,
            max_persona_len, max_context_len, max_response_len,
            max_history_turns, device,
        )

        with torch.no_grad():
            preds = detector.predict(*ids, binary_threshold=binary_threshold)

        sample["d_score"] = preds["d_score"].item()
        sample["c_primary_probs"] = preds["c_primary_probs"].squeeze(0).cpu().numpy().tolist()
        sample["c_primary_idx"] = preds["c_primary"].item()
        sample["e_H_pool"] = preds["e_H_pool"].squeeze(0).cpu().numpy().tolist()
        sample["e_P_pool"] = preds["e_P_pool"].squeeze(0).cpu().numpy().tolist()

        # Keep ground-truth l_risk for reward computation; add detector l_risk separately
        sample["det_l_risk"] = preds["l_risk"].item()

        processed.append(sample)

        if (i + 1) % 100 == 0:
            print(f"Preprocessed {i + 1}/{len(samples)} samples...")

    return processed


def build_obs_vector(sample: Dict, max_turns: int = 20) -> np.ndarray:
    """
    Build the flat observation vector for the RL agent from a preprocessed sample.

    Layout: [d_score(1) | l_risk_onehot(5) | c_primary_probs(10) |
              e_H_pool(H) | e_P_pool(H) | t_norm(1)]
    """
    d_score = np.array([sample["d_score"]], dtype=np.float32)

    l_risk_onehot = np.zeros(NUM_RISK_LEVELS, dtype=np.float32)
    # Use detector-predicted level (det_l_risk), not ground truth, to match deployment
    l_risk_onehot[int(sample.get("det_l_risk", sample["l_risk"]))] = 1.0

    c_primary_probs = np.array(sample["c_primary_probs"], dtype=np.float32)
    e_H_pool = np.array(sample["e_H_pool"], dtype=np.float32)
    e_P_pool = np.array(sample["e_P_pool"], dtype=np.float32)

    num_turns = len(sample.get("history", []))
    t_norm = np.array([min(num_turns / max_turns, 1.0)], dtype=np.float32)

    return np.concatenate([d_score, l_risk_onehot, c_primary_probs, e_H_pool, e_P_pool, t_norm])


def build_bc_tensors(
    processed_samples: List[Dict],
    device: str = "cpu",
) -> tuple:
    """Build (obs_tensor, expert_action_tensor) for behavior cloning warm-up."""
    obs_list, action_list = [], []
    for s in processed_samples:
        obs_list.append(build_obs_vector(s))
        action_list.append(ACTION_NAME_TO_ID.get(s.get("a_recommend", "PASS"), 0))

    obs_tensor = torch.FloatTensor(np.stack(obs_list)).to(device)
    action_tensor = torch.LongTensor(action_list).to(device)
    return obs_tensor, action_tensor
