"""Native GLiNER2 span-architecture schema preprocessing for a fixed Core ML bucket."""
import numpy as np
from gliner2 import Schema
from gliner2.models.base import load_extractor_tokenizer
from gliner2.processor import SchemaTransformer
from gliner2.training.trainer import ExtractorCollator


def load_processor(tokenizer_dir: str):
    """Load only the tokenizer and schema formatter needed by the Core ML model."""
    return SchemaTransformer(tokenizer=load_extractor_tokenizer(tokenizer_dir), token_pooling="first")


def classification_schema(tasks: dict) -> Schema:
    """Same task-dict handling as native ``classify_text``."""
    schema = Schema()
    for name, config in tasks.items():
        if isinstance(config, dict) and "labels" in config:
            cfg = config.copy()
            labels = cfg.pop("labels")
            schema.classification(name, labels, **cfg)
        else:
            schema.classification(name, config)
    return schema


def task_labels(tasks: dict) -> dict[str, list[str]]:
    result = {}
    for name, config in tasks.items():
        labels = config["labels"] if isinstance(config, dict) and "labels" in config else config
        result[name] = list(labels.keys()) if isinstance(labels, dict) else list(labels)
    return result


def prepare_decision(processor, text: str, tasks: dict, length: int, max_heads: int, max_options: int):
    """Tokenize ``tasks`` exactly as the native span collator does and pad into the bucket."""
    if not 1 <= len(tasks) <= max_heads:
        raise ValueError(f"Expected 1..{max_heads} decision heads, got {len(tasks)}")
    labels = task_labels(tasks)
    for name, values in labels.items():
        if not 1 <= len(values) <= max_options:
            raise ValueError(f"Head {name!r} has {len(values)} labels; bucket holds 1..{max_options}")
    collator = ExtractorCollator(processor, is_training=False, max_len=None, architecture="span")
    batch = collator([(text, classification_schema(tasks).build())])
    ids = batch.input_ids.numpy()
    if ids.shape[1] > length:
        raise ValueError(f"Schema and text require {ids.shape[1]} subwords; bucket holds {length}")
    groups = batch.schema_special_indices[0]
    if len(groups) != len(labels):
        raise ValueError("Schema head count does not match the requested tasks")
    indices = np.zeros((1, max_heads, max_options), dtype=np.int32)
    mask = np.zeros((1, max_heads, max_options), dtype=np.float32)
    for head, (positions, values) in enumerate(zip(groups, labels.values())):
        # positions[0] is the [P] prompt marker; the rest are one [L] marker per label.
        markers = list(positions[1:])
        if len(markers) != len(values):
            raise ValueError("Label markers were truncated or merged")
        indices[0, head, : len(markers)] = markers
        mask[0, head, : len(markers)] = 1.0
    attention = batch.attention_mask.numpy()
    pad = processor.tokenizer.pad_token_id
    return {
        "input_ids": np.pad(ids, ((0, 0), (0, length - ids.shape[1])), constant_values=pad).astype(np.int32),
        "attention_mask": np.pad(attention, ((0, 0), (0, length - attention.shape[1]))).astype(np.int32),
        "marker_indices": indices,
        "marker_mask": mask,
    }