import json
from transformers import AutoTokenizer
import torch


class StepPreprocessor:
    def __init__(self, model_name="microsoft/deberta-v3-base", max_len=512):
        # use_fast=False is required for DeBERTa-v3 to avoid spm.model parsing errors on Linux
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
        self.max_len = max_len

    def encode_step(self, query, step):
        """
        Encodes a single step as a concatenated string:
          [QUERY] ... [THOUGHT] ... [ACTION] ... [OBS] ...

        This is the "single-channel concatenated" representation where the
        model learns to attend across all information jointly.
        """
        # Channel 1 — Reasoning trace (the LLM's textual thought / plan)
        thought = step.get("content", "") or ""

        # Channel 2 — Tool context (action + observation)
        tool_calls = step.get("tool_calls") or []
        tool_responses = step.get("tool_responses") or []

        actions_str = ""
        for tc in tool_calls:
            name = tc.get("name", "")
            args = json.dumps(tc.get("arguments", {}))
            actions_str += f"{name} {args} | "

        observations_str = " | ".join(str(tr) for tr in tool_responses) if tool_responses else ""

        tool_context = f"[ACTION] {actions_str.strip()} [OBS] {observations_str.strip()}"

        # Channel 3 — Query context (the original question / task)
        if isinstance(query, list):
            flat_query = " ".join(q[0] if isinstance(q, list) else str(q) for q in query)
        else:
            flat_query = str(query) if query else ""

        # Truncation strategy: keep the query short but preserve tool context
        # because that is the dominant signal (ablation: -28.8 pt without it).
        # We truncate THOUGHT in the middle if necessary.
        full_text = f"[QUERY] {flat_query} [THOUGHT] {thought} {tool_context}"

        encoded = self.tokenizer(
            full_text,
            max_length=self.max_len,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )
        return encoded

    def encode_trajectory(self, sample):
        """
        Returns a list of dicts, one per step, each containing:
          - encoding: tokenized input (input_ids + attention_mask)
          - label:    1 if this step is the hallucination step, else 0
          - step_idx: the integer step index from the JSON
        """
        # Support both 'trajectory' and 'history' JSON keys
        trajectory = sample.get("trajectory") or sample.get("history") or []

        is_hal = sample.get("is_hallucination", False)
        if isinstance(is_hal, str):
            is_hal = is_hal.lower() == "true"

        hal_step = sample.get("hallucination_step")
        if hal_step is not None:
            hal_step = int(hal_step)

        query = sample.get("question", sample.get("query", ""))

        steps = []
        for step in trajectory:
            enc = self.encode_step(query, step)
            step_idx = step.get("step")

            # A step is labelled positive only if the trajectory is hallucinated
            # AND this is exactly the annotated root-cause step.
            label = 1 if (is_hal and step_idx == hal_step) else 0

            steps.append({
                "encoding": enc,
                "label": label,
                "step_idx": step_idx,
            })

        return steps


# ── quick smoke-test ─────────────────────────────────────────────────────────
if __name__ == "__main__":
    import os

    print("Testing StepPreprocessor (will download tokenizer if not cached)…")
    pre = StepPreprocessor()

    script_dir = os.path.dirname(os.path.abspath(__file__))
    sample_file = os.path.join(script_dir, "..", "..", "data", "splits", "train.json")

    if os.path.exists(sample_file):
        with open(sample_file, "r", encoding="utf-8") as fh:
            samples = json.load(fh)
        sample = samples[0]
        encoded = pre.encode_trajectory(sample)
        print(f"OK — encoded {len(encoded)} steps.")
        print(f"input_ids shape: {encoded[0]['encoding']['input_ids'].shape}")
    else:
        print(f"File not found: {sample_file}")