import json
import re
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional
from importlib import resources

import numpy as np
import onnxruntime as ort

KHMER_RE = re.compile(r"[\u1780-\u17FF\u19E0-\u19FF]+")


def _read_json(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def _default_resource_path(name: str) -> Path:
    return Path(resources.files("khmersm").joinpath(name))


def _encode_chars(text: str, char2id: dict) -> List[int]:
    unk = char2id.get("<unk>", 1)
    return [char2id.get(ch, unk) for ch in text]


def _viterbi_decode(emissions: np.ndarray, transitions: np.ndarray, start: np.ndarray, end: np.ndarray) -> List[int]:
    # emissions: (T, C)
    T, C = emissions.shape
    score = start + emissions[0]
    backpointers = []

    for t in range(1, T):
        prev_score = score.reshape(C, 1) + transitions  # (C, C)
        best_prev = prev_score.argmax(axis=0)
        score = prev_score.max(axis=0) + emissions[t]
        backpointers.append(best_prev)

    score = score + end
    best_last = int(score.argmax(axis=0))

    tags = [best_last]
    for bp in reversed(backpointers):
        best_last = int(bp[best_last])
        tags.append(best_last)
    tags.reverse()
    return tags


def _segment_with_tags(text: str, tag_ids: List[int], id2tag: dict) -> str:
    out = []
    for i, ch in enumerate(text):
        tag = id2tag[str(tag_ids[i])] if isinstance(id2tag, dict) else id2tag[tag_ids[i]]
        if i > 0 and tag == "B":
            out.append(" ")
        out.append(ch)
    return "".join(out)


@dataclass
class Segmenter:
    config_path: Optional[Path] = None
    onnx_path: Optional[Path] = None

    def __post_init__(self) -> None:
        if self.config_path is None:
            self.config_path = _default_resource_path("config.json")
        if self.onnx_path is None:
            self.onnx_path = _default_resource_path("model.onnx")

        if not self.config_path.exists():
            raise FileNotFoundError(f"Missing config: {self.config_path}")
        if not self.onnx_path.exists():
            raise FileNotFoundError(f"Missing ONNX model: {self.onnx_path}")

        self.config = _read_json(self.config_path)
        self.char2id = self.config["char2id"]
        self.tag2id = self.config["tag2id"]
        self.id2tag = {str(v): k for k, v in self.tag2id.items()}

        self.transitions = np.array(self.config["crf_transitions"], dtype=np.float32)
        self.start = np.array(self.config["crf_start_transitions"], dtype=np.float32)
        self.end = np.array(self.config["crf_end_transitions"], dtype=np.float32)

        self.sess = ort.InferenceSession(str(self.onnx_path), providers=["CPUExecutionProvider"])

    def _segment_khmer(self, segment: str) -> str:
        x_ids = _encode_chars(segment, self.char2id)
        x = np.array([x_ids], dtype=np.int64)
        mask = np.ones_like(x, dtype=bool)
        emissions = self.sess.run(None, {"x": x, "mask": mask})[0][0]
        tags = _viterbi_decode(emissions, self.transitions, self.start, self.end)
        return _segment_with_tags(segment, tags, self.id2tag)

    def segment(self, text: str) -> str:
        if not text:
            return ""

        parts = []
        last = 0
        for m in KHMER_RE.finditer(text):
            if m.start() > last:
                parts.append(text[last:m.start()])
            parts.append(self._segment_khmer(m.group(0)))
            last = m.end()
        if last < len(text):
            parts.append(text[last:])

        return "".join(parts)


def segment(text: str) -> str:
    return Segmenter().segment(text)
