"""
Debatra — Worker 1: ADU Parser Inference
==========================================
Base model:  microsoft/deberta-v3-base
Adapter:     LoRA r=16 trained on climate_fever + DebateSum
Task:        BIO token classification
Labels:      O=0  B-claim=1  I-claim=2  B-warrant=3  I-warrant=4
Test F1:     0.836

Input:  raw text string
Output: {
  "tokens":       ["word1", "word2", ...],
  "labels":       [0, 1, 2, ...],          # BIO int per token
  "segments": [
    {"type": "claim",   "text": "..."},
    {"type": "warrant", "text": "..."},
  ],
  "structure":    "claim → warrant",        # human readable
  "completeness": 0.82,                     # 0-1 score
  "score":        8.2,                      # 0-10 for composite
  "uncertain":    False,
}
"""

import re
import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification
from peft import PeftModel


LABEL2ID = {"O":0, "B-claim":1, "I-claim":2, "B-warrant":3, "I-warrant":4}
ID2LABEL  = {v: k for k, v in LABEL2ID.items()}
BASE_MODEL = "microsoft/deberta-v3-base"
MAX_LENGTH = 128


class ADUParserWorker:
    def __init__(self, model_path: str, confidence_threshold: float = 0.70):
        self.model_path           = model_path
        self.confidence_threshold = confidence_threshold
        self._loaded              = False
        self._load()

    def _load(self):
        try:
            try:
                self.tokenizer = AutoTokenizer.from_pretrained(
                    self.model_path,
                    use_fast=True,
                )
            except Exception:
                # Some LoRA export folders omit config.json; fall back to base tokenizer.
                self.tokenizer = AutoTokenizer.from_pretrained(
                    BASE_MODEL,
                    use_fast=True,
                )
            base = AutoModelForTokenClassification.from_pretrained(
                BASE_MODEL,
                num_labels=len(LABEL2ID),
                ignore_mismatched_sizes=True,
            )
            self.model = PeftModel.from_pretrained(base, self.model_path)
            self.model.eval()
            if torch.cuda.is_available():
                self.model = self.model.cuda()
            self._loaded = True
        except Exception as e:
            raise RuntimeError(f"ADU Parser failed to load from {self.model_path}: {e}")

    def status(self) -> str:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        return f"loaded ({device})" if self._loaded else "not loaded"

    def predict(self, text: str) -> dict:
        """Run BIO inference on input text."""
        tokens = text.split()
        if not tokens:
            return self._empty_result()

        enc = self.tokenizer(
            tokens,
            is_split_into_words=True,
            truncation=True,
            max_length=MAX_LENGTH,
            return_tensors="pt",
        )

        device = next(self.model.parameters()).device
        enc    = {k: v.to(device) for k, v in enc.items()}

        with torch.no_grad():
            outputs = self.model(**enc)
            logits  = outputs.logits[0]
            probs   = torch.softmax(logits, dim=-1)
            preds   = torch.argmax(probs, dim=-1)
            confs   = probs.max(dim=-1).values

        # Align predictions back to word level
        word_ids     = enc["input_ids"][0] if "input_ids" in enc else []
        enc_word_ids = self.tokenizer(
            tokens,
            is_split_into_words=True,
            truncation=True,
            max_length=MAX_LENGTH,
        ).word_ids()

        word_labels = {}
        word_confs  = {}
        for idx, word_id in enumerate(enc_word_ids):
            if word_id is None:
                continue
            if word_id not in word_labels:
                word_labels[word_id] = preds[idx].item()
                word_confs[word_id]  = confs[idx].item()

        predicted_labels = [word_labels.get(i, 0) for i in range(len(tokens))]
        avg_conf         = sum(word_confs.values()) / max(len(word_confs), 1)

        # Extract segments
        segments    = self._extract_segments(tokens, predicted_labels)
        structure   = self._describe_structure(segments)
        completeness= self._compute_completeness(segments)

        uncertain = avg_conf < self.confidence_threshold
        score     = round(completeness * 10, 2)

        return {
            "tokens":       tokens,
            "labels":       predicted_labels,
            "segments":     segments,
            "structure":    structure,
            "completeness": completeness,
            "score":        score,
            "confidence":   round(avg_conf, 3),
            "uncertain":    uncertain,
        }

    def _extract_segments(self, tokens, labels):
        segments = []
        cur_type = None
        cur_words= []
        for word, label_id in zip(tokens, labels):
            label = ID2LABEL[label_id]
            if label.startswith("B-"):
                if cur_type and cur_words:
                    segments.append({"type": cur_type, "text": " ".join(cur_words)})
                cur_type  = label[2:]   # "claim" or "warrant"
                cur_words = [word]
            elif label.startswith("I-") and cur_type:
                cur_words.append(word)
            else:  # O
                if cur_type and cur_words:
                    segments.append({"type": cur_type, "text": " ".join(cur_words)})
                cur_type  = None
                cur_words = []
        if cur_type and cur_words:
            segments.append({"type": cur_type, "text": " ".join(cur_words)})
        return segments

    def _describe_structure(self, segments):
        if not segments:
            return "no_structure"
        types = [s["type"] for s in segments]
        return " → ".join(types)

    def _compute_completeness(self, segments):
        """
        Score: claim present = 0.4, warrant present = 0.4, both = 0.8+
        Multiple warrants = bonus up to 1.0.
        """
        has_claim   = any(s["type"] == "claim"   for s in segments)
        n_warrants  = sum(1 for s in segments if s["type"] == "warrant")
        score = 0.0
        if has_claim:   score += 0.4
        if n_warrants:  score += min(0.4 + (n_warrants - 1) * 0.1, 0.6)
        return round(min(score, 1.0), 3)

    def _empty_result(self):
        return {
            "tokens":[], "labels":[], "segments":[],
            "structure":"no_structure", "completeness":0.0,
            "score":0.0, "confidence":0.0, "uncertain":True,
        }