"""
Batch annotation for the Annotate & Export tab and its API endpoints.
"""

import re
from typing import Dict, List, Optional

from core.treebanks_data import AVAILABLE_MODELS
from core.udpipe_client import UDPipeClient

MAX_LINES = 50
_DOCUMENT_COMMENTS = ("# generator", "# udpipe_model", "# newdoc", "# newpar")


def model_ids() -> List[str]:
    """All selectable UDPipe model IDs, English first."""
    return [model_id for models in AVAILABLE_MODELS.values() for model_id in models.values()]


def resolve_model(model: str) -> str:
    """Accepts a model ID or a UI label such as 'UDPipe · german-gsd-ud-2.17-251125'."""
    for models in AVAILABLE_MODELS.values():
        for label, model_id in models.items():
            if model in (label, model_id):
                return model_id
    raise ValueError(f"Unknown model {model!r}. Available: {', '.join(model_ids())}")


def merge_conllu(parts: List[str], model_id: str) -> str:
    """
    Joins per-line parser outputs into one document with one header and
    consecutive sent_id values. Sentence comments and token lines are kept.
    """
    sentences = []
    for part in parts:
        for block in re.split(r"\n\s*\n", part.strip()):
            lines = [line for line in block.splitlines()
                     if not line.startswith(_DOCUMENT_COMMENTS) and line.strip()]
            if any(not line.startswith("#") for line in lines):
                sentences.append(lines)
    header = ["# generator = UDPipe 2, https://lindat.mff.cuni.cz/services/udpipe",
              f"# udpipe_model = {model_id}"]
    blocks = []
    for number, lines in enumerate(sentences, start=1):
        lines = [line for line in lines if not line.startswith("# sent_id")]
        blocks.append("\n".join([f"# sent_id = {number}"] + lines))
    return "\n".join(header) + "\n" + "\n\n".join(blocks) + "\n\n"


def annotate_text(
    text: str,
    model: str,
    include_morphology: bool = False,
    client: Optional[UDPipeClient] = None,
) -> Dict[str, object]:
    """
    Parses each non-empty input line separately and returns
    {"conllu": str, "model": str, "sentences": int}. Raises ValueError on bad input
    and RuntimeError when the parser service fails.
    """
    model_id = resolve_model(model)
    lines = [line.strip() for line in (text or "").splitlines() if line.strip()]
    if not lines:
        raise ValueError("Enter at least one sentence.")
    if len(lines) > MAX_LINES:
        raise ValueError(f"At most {MAX_LINES} lines per request.")
    client = client or UDPipeClient()
    parts = []
    for line in lines:
        result = client.parse_sentence(line, model_id, strip_morphology=not include_morphology)
        if result.get("error"):
            raise RuntimeError(result["error"])
        parts.append(result["conllu"])
    conllu = merge_conllu(parts, model_id)
    return {"conllu": conllu, "model": model_id, "sentences": conllu.count("# sent_id")}