"""
Client wrapper for LINDAT UDPipe 2 REST API.
API Reference: https://lindat.mff.cuni.cz/services/udpipe/api/
"""

import requests
from typing import Dict, Any, Optional

UDPIPE_API_URL = "https://lindat.mff.cuni.cz/services/udpipe/api/process"


class UDPipeClient:
    def __init__(self, timeout: int = 15):
        self.timeout = timeout
        self.session = requests.Session()

    def parse_sentence(
        self,
        text: str,
        model: str,
        strip_morphology: bool = True,
        cache: Optional[Any] = None
    ) -> Dict[str, Any]:
        """
        Parses a single sentence with UDPipe 2.
        
        Args:
            text: Sentence text (single sentence).
            model: Model identifier (e.g. 'english-ewt-ud-2.17-251125').
            strip_morphology: If True, sets FEATS column to '_' so the visualizer
                              and parser focus strictly on POS tags and dependencies.
            cache: Optional BoundedParseCache or mapping with .get(model, text) and
                   .set(model, text, conllu).
        
        Returns:
            Dictionary with { "conllu": str, "model": str, "error": Optional[str] }
        """
        clean_text = text.strip()
        # Enforce single sentence processing (first sentence line)
        if "\n" in clean_text:
            clean_text = clean_text.split("\n")[0].strip()

        if not clean_text:
            return {"conllu": "", "model": model, "error": "Empty input text."}

        if cache is not None:
            cached_conllu = cache.get(model, clean_text, strip_morphology=strip_morphology)
            if cached_conllu is not None:
                return {
                    "conllu": cached_conllu,
                    "model": model,
                    "error": None
                }

        data = {
            "data": clean_text,
            "model": model,
            "tokenizer": "",
            "tagger": "",
            "parser": "",
        }

        try:
            resp = self.session.post(UDPIPE_API_URL, data=data, timeout=self.timeout)
            resp.raise_for_status()
            res_json = resp.json()
            conllu_raw = res_json.get("result", "")

            if strip_morphology:
                conllu_processed = self._strip_feats_column(conllu_raw)
            else:
                conllu_processed = conllu_raw

            if cache is not None and conllu_processed:
                cache.set(model, clean_text, conllu_processed, strip_morphology=strip_morphology)

            return {
                "conllu": conllu_processed,
                "model": model,
                "error": None
            }
        except requests.exceptions.RequestException as e:
            return {
                "conllu": "",
                "model": model,
                "error": f"UDPipe API Request Error: {e}"
            }

    @staticmethod
    def _strip_feats_column(conllu_text: str) -> str:
        """
        Replaces column 6 (FEATS) with '_' to keep focus on POS tags and dependency relations.
        CoNLL-U columns: 1:ID, 2:FORM, 3:LEMMA, 4:UPOS, 5:XPOS, 6:FEATS, 7:HEAD, 8:DEPREL, 9:DEPS, 10:MISC
        """
        out_lines = []
        for line in conllu_text.splitlines():
            line_str = line.strip()
            if not line_str or line_str.startswith("#"):
                out_lines.append(line)
                continue
            cols = line.split("\t")
            if len(cols) == 10:
                cols[5] = "_"  # Blank morphological features
                out_lines.append("\t".join(cols))
            else:
                out_lines.append(line)
        return "\n".join(out_lines) + "\n"