"""词典导入共享工具。"""

from __future__ import annotations

import re
from typing import Optional

from sqlalchemy.dialects.mysql import insert as mysql_insert
from sqlalchemy.orm import Session

from models import DictionaryEntry

_POS_PREFIX = re.compile(r"^[a-zA-Z]+\.\s*")
_TAG_PREFIX = re.compile(r"^\[[^\]]+\]\s*")


def clean_translation(text: str, max_len: int = 256) -> str:
    if not text:
        return ""
    lines = [ln.strip() for ln in text.replace("\r", "").split("\n") if ln.strip()]
    parts: list[str] = []
    for ln in lines:
        ln = _TAG_PREFIX.sub("", ln)
        ln = _POS_PREFIX.sub("", ln)
        if ln:
            parts.append(ln)
    result = "；".join(parts) if parts else text.strip()
    return result[:max_len]


def normalize_row(raw: dict, source: Optional[str]) -> Optional[dict]:
    lemma_en = (raw.get("lemma_en") or raw.get("en") or raw.get("word") or "").strip().lower()
    zh = clean_translation(raw.get("zh") or raw.get("cn") or raw.get("translation") or "")
    if not lemma_en or not zh:
        return None
    if len(lemma_en) > 128:
        return None
    phonetic = (raw.get("phonetic") or "").strip() or None
    if phonetic and len(phonetic) > 128:
        phonetic = phonetic[:128]
    return {
        "lemma_en": lemma_en,
        "zh": zh,
        "phonetic": phonetic,
        "example_en": (raw.get("example_en") or "").strip() or None,
        "example_cn": (raw.get("example_cn") or "").strip() or None,
        "source": source or (raw.get("source") or "import"),
    }


def upsert_batch(session: Session, rows: list[dict]) -> int:
    if not rows:
        return 0
    stmt = mysql_insert(DictionaryEntry).values(rows)
    stmt = stmt.on_duplicate_key_update(
        zh=stmt.inserted.zh,
        phonetic=stmt.inserted.phonetic,
        example_en=stmt.inserted.example_en,
        example_cn=stmt.inserted.example_cn,
        source=stmt.inserted.source,
    )
    session.execute(stmt)
    return len(rows)


def upsert_entries(session: Session, raw_rows: list[dict], source: Optional[str], batch_size: int = 2000) -> int:
    batch: list[dict] = []
    total = 0
    for raw in raw_rows:
        row = normalize_row(raw, source)
        if not row:
            continue
        batch.append(row)
        if len(batch) >= batch_size:
            upsert_batch(session, batch)
            session.commit()
            total += len(batch)
            batch.clear()
    if batch:
        upsert_batch(session, batch)
        session.commit()
        total += len(batch)
    return total
