"""BFB232 prototype inference. MIT license; see LICENSE and THIRD_PARTY_NOTICES.md."""
from __future__ import annotations

import argparse
import bisect
import contextlib
import json
import math
import os
from pathlib import Path
import threading
import time

# Set before the first CUDA operation. Matches the reference experiment.
os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")

import torch
from safetensors.torch import load_file
from transformers import AutoModel, AutoTokenizer

import bfb232_text as textio


class BoundaryHead(torch.nn.Module):
    """Frozen B187/B232 'left' binary head; CLOSE score = logit[1] - logit[0]."""

    def __init__(self):
        super().__init__()
        self.base = torch.nn.Sequential(
            torch.nn.LayerNorm(768), torch.nn.Linear(768, 96),
            torch.nn.GELU(), torch.nn.Linear(96, 2),
        )
        self.extra = torch.nn.Sequential(
            torch.nn.LayerNorm(384), torch.nn.Linear(384, 64),
            torch.nn.GELU(), torch.nn.Linear(64, 2),
        )

    def forward(self, features):
        return self.base(features[:, :768]) + self.extra(features[:, 768:1152])


def make_window(tokenizer, text, positions, left, right):
    """Exact bfb130 indexed token anchors and 256-character side pools."""
    encoding = tokenizer(text[left:right], return_offsets_mapping=True,
                         add_special_tokens=True, truncation=False)
    offsets = encoding.pop("offset_mapping")
    if len(encoding["input_ids"]) > 8192:
        raise ValueError("Retokenized window exceeds the reference 8192-token limit")
    valid = [i for i, (a, b) in enumerate(offsets) if b > a]
    starts = [offsets[i][0] for i in valid]
    if starts != sorted(starts):
        raise ValueError("Non-monotone tokenizer offsets")
    anchors, after, sides = [], [], []
    for position in positions:
        pos = position - left
        at = bisect.bisect_left(starts, pos)
        if at == 0:
            raise ValueError("Candidate has no preceding token")
        a = valid[at - 1]
        b = valid[at] if at < len(valid) else a
        lo = bisect.bisect_left(starts, max(0, pos - 256))
        hi = bisect.bisect_left(starts, pos + 256)
        anchors.append(a)
        after.append(b)
        sides.append((valid[lo:at] or [a], valid[at:hi] or [b]))
    return dict(encoding=dict(encoding), anchors=anchors, after=after,
                sides=sides, left=left, right=right)


def inference_windows(tokenizer, text, positions):
    """Canonical text/positions; 7900-token lattice, stride 6500, interior owner."""
    encoded = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True,
                        truncation=False, verbose=False)
    offsets = encoded["offset_mapping"]
    if not positions:
        return [], len(offsets)
    if len(offsets) <= 7900:
        w = make_window(tokenizer, text, positions, 0, len(text))
        w["indices"] = list(range(len(positions)))
        return [w], len(offsets)
    starts = sorted(set(list(range(0, len(offsets) - 7900 + 1, 6500))
                        + [len(offsets) - 7900]))
    spans = [(offsets[s][0] if s else 0,
              offsets[s + 7899][1] if s + 7900 < len(offsets) else len(text))
             for s in starts]
    groups = {}
    for i, p in enumerate(positions):
        choices = [j for j, (a, b) in enumerate(spans) if a < p <= b]
        if not choices:
            raise ValueError("Candidate lies outside the token windows")
        owner = max(choices, key=lambda j: min(p - spans[j][0], spans[j][1] - p))
        groups.setdefault(owner, []).append(i)
    result = []
    for j, indices in sorted(groups.items()):
        w = make_window(tokenizer, text, [positions[i] for i in indices], *spans[j])
        w["indices"] = indices
        result.append(w)
    return result, len(offsets)


def features_from_hidden(hidden, window):
    """Preserves reference FP32 pooling and operation order."""
    pair = torch.cat([hidden[window["anchors"]], hidden[window["after"]]], -1).float()
    left = torch.stack([hidden[a].float().mean(0) for a, _ in window["sides"]])
    right = torch.stack([hidden[b].float().mean(0) for _, b in window["sides"]])
    return torch.cat([pair, left, right, right - left, right * left], -1)


class BFB232:
    def __init__(self, tokenizer, encoder, head, config, device):
        self.tokenizer, self.encoder, self.head = tokenizer, encoder, head
        self.config, self.device = config, torch.device(device)
        self.thresholds = dict(config["thresholds"])
        self.lock = threading.Lock()

    @classmethod
    def from_pretrained(cls, model_path, *, device="auto", revision=None,
                        local_files_only=False, token=None):
        """Load a local release folder or Hub ID; never executes Hub Python code.

        For a Hub ID, pin revision to its commit hash. CPU uses FP32; the
        reference CUDA path uses BF16 autocast with FP32 stored weights.
        """
        folder = Path(model_path)
        if not folder.is_dir():
            from huggingface_hub import snapshot_download
            folder = Path(snapshot_download(
                repo_id=str(model_path), revision=revision, token=token,
                local_files_only=local_files_only,
                allow_patterns=["bfb232_config.json", "encoder/*.json",
                                "encoder/*.safetensors", "tokenizer/*.json",
                                "boundary_head.safetensors"],
            ))
        config = json.loads((folder / "bfb232_config.json").read_text("utf-8"))
        if config["format_version"] != 1 or config["input_contract"] != textio.VERSION:
            raise ValueError("Unsupported BFB232 release/input contract")
        if device == "auto":
            device = "cuda" if torch.cuda.is_available() and torch.cuda.is_bf16_supported(including_emulation=False) else "cpu"
        device = torch.device(device)
        if device.type not in ("cpu", "cuda"):
            raise ValueError("Only CPU FP32 and CUDA BF16 are supported")
        if device.type == "cuda":
            if not torch.cuda.is_available():
                raise ValueError("CUDA is unavailable; use device='cpu'")
            with torch.cuda.device(device):
                if not torch.cuda.is_bf16_supported(including_emulation=False):
                    raise ValueError("CUDA BF16 support is required; use device='cpu'")
        torch.set_num_threads(4)
        torch.use_deterministic_algorithms(True)
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False
        tokenizer = AutoTokenizer.from_pretrained(folder / "tokenizer", local_files_only=True)
        encoder, info = AutoModel.from_pretrained(
            folder / "encoder", local_files_only=True, use_safetensors=True,
            dtype=torch.float32, attn_implementation="sdpa", output_loading_info=True,
        )
        if any(info.get(k) for k in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")):
            raise ValueError(f"Encoder state mismatch: {info}")
        head = BoundaryHead()
        head.load_state_dict(load_file(str(folder / "boundary_head.safetensors")), strict=True)
        encoder.to(device).eval().requires_grad_(False)
        head.to(device).eval().requires_grad_(False)
        return cls(tokenizer, encoder, head, config, device)

    def segment(self, text, language):
        """Return source-preserving segments and original Unicode code-point offsets.

        Language must be 'ko' or 'en'; no automatic language detection or
        minimum/maximum segment length is imposed. Empty input returns [].
        """
        if not isinstance(text, str):
            raise TypeError("text must be a string")
        if language not in self.thresholds:
            raise ValueError("language must be 'ko' or 'en'")
        with self.lock:
            return self._segment(text, language)

    def _segment(self, text, language):
        started = time.perf_counter()
        canonical = textio.normalize(text)
        positions = textio.candidates(canonical.text, language)
        original_positions = [canonical.to_original[p] for p in positions]
        windows, input_tokens = inference_windows(self.tokenizer, canonical.text, positions)
        scores = [None] * len(positions)
        processed_tokens = 0
        autocast = (torch.autocast("cuda", dtype=torch.bfloat16)
                    if self.device.type == "cuda" else contextlib.nullcontext())
        with torch.inference_mode(), autocast:
            for w in windows:
                inputs = self.tokenizer.pad([w["encoding"]], padding=True, return_tensors="pt")
                inputs = {k: v.to(self.device) for k, v in inputs.items()}
                hidden = self.encoder(**inputs).last_hidden_state[0]
                logits = self.head(features_from_hidden(hidden, w))
                values = (logits[:, 1] - logits[:, 0]).float().cpu().tolist()
                for i, score in zip(w["indices"], values):
                    if scores[i] is not None or not math.isfinite(score):
                        raise RuntimeError("Invalid/duplicate boundary score")
                    scores[i] = score
                processed_tokens += len(w["encoding"]["input_ids"])
        if any(s is None for s in scores):
            raise RuntimeError("Missing boundary scores")
        cuts = [p for p, s in zip(original_positions, scores) if s >= self.thresholds[language]]
        segments, left = [], 0
        for right in cuts + ([len(text)] if text else []):
            segments.append(dict(index=len(segments) + 1, start=left, end=right,
                                 text=text[left:right], boundary="END" if right == len(text) else "CLOSE"))
            left = right
        if "".join(s["text"] for s in segments) != text:
            raise RuntimeError("Source preservation failed")
        return dict(model="BFB232", language=language, segments=segments,
                    source_preserved=True, input_contract=textio.VERSION,
                    candidate_positions=original_positions, scores=scores,
                    threshold=self.thresholds[language],
                    precision="cuda-bf16" if self.device.type == "cuda" else "cpu-fp32",
                    metrics=dict(input_chars=len(text), input_tokens=input_tokens,
                                 processed_tokens=processed_tokens, windows=len(windows),
                                 candidates=len(positions), segment_count=len(segments),
                                 total_seconds=time.perf_counter() - started))


def main():
    parser = argparse.ArgumentParser(description="BFB232 semantic segmentation")
    parser.add_argument("--model", default=str(Path(__file__).parent))
    parser.add_argument("--revision", help="Hub commit hash when --model is a Hub ID")
    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"])
    parser.add_argument("--language", required=True, choices=["ko", "en"])
    parser.add_argument("--input", required=True, type=Path, help="UTF-8 text file")
    parser.add_argument("--output", required=True, type=Path, help="New JSON file; refuses overwrite")
    args = parser.parse_args()
    if args.output.exists():
        parser.error("Output exists; choose a new path")
    # newline='' preserves CRLF/CR in the returned original-text slices.
    with args.input.open(encoding="utf-8", newline="") as source:
        text = source.read()
    model = BFB232.from_pretrained(args.model, device=args.device, revision=args.revision)
    result = model.segment(text, args.language)
    with args.output.open("x", encoding="utf-8", newline="\n") as out:
        json.dump(result, out, ensure_ascii=False, indent=2)
        out.write("\n")


if __name__ == "__main__":
    main()