smoke_signal/scripts/font_library.py

"""
Smoke Signal — Font Detection + Font-Aware OCR Helpers
======================================================
Lightweight utilities for:
1) identifying page fonts via MixFont Lens API
2) maintaining per-font OCR preference metadata
3) selecting per-font custom Tesseract models when available
4) exposing font-specific punctuation-map paths

This module is intentionally defensive:
- if MixFont key is missing, it degrades gracefully
- if no custom .traineddata exists, it falls back to "eng"
"""

from __future__ import annotations

import json
import os
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Optional

SS_ROOT = Path(os.environ.get("SS_DATA_ROOT", os.environ.get("SS_ROOT", "/tmp/smoke_signal")))
CALIBRATION_DIR = SS_ROOT / "calibration"
FONTS_DIR = SS_ROOT / "fonts"
TESSDATA_DIR = SS_ROOT / "tessdata"
REPO_TESSDATA_DIR = Path(__file__).resolve().parents[1] / "tessdata"
FONT_REGISTRY_PATH = CALIBRATION_DIR / "font_registry.json"

for _dir in (CALIBRATION_DIR, FONTS_DIR, TESSDATA_DIR, REPO_TESSDATA_DIR):
    _dir.mkdir(parents=True, exist_ok=True)

MIXFONT_API_KEY = os.environ.get("MIXFONT_API_KEY", "").strip()
# MixFont Lens API docs endpoint.
MIXFONT_API_URL = os.environ.get("MIXFONT_API_URL", "https://www.mixfont.com/v1/api/lens").strip()
try:
    MIXFONT_TIMEOUT_SEC = max(1.0, float(os.environ.get("MIXFONT_TIMEOUT_SEC", "4")))
except Exception:
    MIXFONT_TIMEOUT_SEC = 4.0


def _slug_font_name(font_name: str) -> str:
    cleaned = "".join(ch.lower() if ch.isalnum() else "_" for ch in (font_name or "").strip())
    while "__" in cleaned:
        cleaned = cleaned.replace("__", "_")
    return cleaned.strip("_")


def _load_registry() -> dict:
    if not FONT_REGISTRY_PATH.exists():
        return {}
    try:
        return json.loads(FONT_REGISTRY_PATH.read_text(encoding="utf-8"))
    except Exception:
        return {}


def _save_registry(registry: dict) -> None:
    FONT_REGISTRY_PATH.write_text(
        json.dumps(registry, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )


def _ensure_registry_entry(font_name: str) -> dict:
    registry = _load_registry()
    entry = registry.get(font_name, {})
    if not entry:
        entry = {
            "identified_count": 0,
            "status": "unknown",  # unknown | available | missing | specimen_only
            "preferred_engine": "surya",  # surya | tesseract
            "tesseract_model": None,
            "avg_confidence_surya": None,
            "avg_confidence_tess": None,
            "last_identified_at": None,
        }
        registry[font_name] = entry
        _save_registry(registry)
    return entry


def _build_public_image_url(local_image_path: str) -> Optional[str]:
    """
    Construct a publicly reachable URL for a local rendered page image.
    Priority:
      1) MIXFONT_IMAGE_URL_TEMPLATE with {path} placeholder
      2) MIXFONT_IMAGE_BASE_URL + /gradio_api/file=<abs_path>
      3) SPACE_HOST + /gradio_api/file=<abs_path>
      4) HF_SPACE_URL + /gradio_api/file=<abs_path>
    """
    raw_path = str(Path(local_image_path).resolve())
    encoded_path = urllib.parse.quote(raw_path, safe="")

    template = os.environ.get("MIXFONT_IMAGE_URL_TEMPLATE", "").strip()
    if template:
        if "{path}" in template:
            return template.replace("{path}", encoded_path)
        return template

    base = (
        os.environ.get("MIXFONT_IMAGE_BASE_URL", "").strip()
        or os.environ.get("HF_SPACE_URL", "").strip()
        or ""
    )
    if not base:
        space_id = os.environ.get("SPACE_ID", "").strip()
        if space_id:
            base = f"https://huggingface.co/spaces/{space_id}"
    if not base:
        space_host = os.environ.get("SPACE_HOST", "").strip()
        if space_host:
            if space_host.startswith("http://") or space_host.startswith("https://"):
                base = space_host
            else:
                base = f"https://{space_host}"

    if not base:
        return None

    return f"{base.rstrip('/')}/gradio_api/file={encoded_path}"


def mixfont_preflight() -> dict:
    """
    Return runtime readiness checks for MixFont integration.
    """
    template = os.environ.get("MIXFONT_IMAGE_URL_TEMPLATE", "").strip()
    base = (
        os.environ.get("MIXFONT_IMAGE_BASE_URL", "").strip()
        or os.environ.get("HF_SPACE_URL", "").strip()
        or ""
    )
    if not base:
        space_id = os.environ.get("SPACE_ID", "").strip()
        if space_id:
            base = f"https://huggingface.co/spaces/{space_id}"
    space_host = os.environ.get("SPACE_HOST", "").strip()
    has_public_source = bool(template or base or space_host)
    return {
        "api_key_set": bool(MIXFONT_API_KEY),
        "api_url": MIXFONT_API_URL,
        "image_url_template_set": bool(template),
        "image_base_set": bool(base),
        "space_host_set": bool(space_host),
        "public_image_url_source_available": has_public_source,
    }


def identify_page_font(image_path: str, image_url: Optional[str] = None) -> dict:
    """
    Identify font via MixFont Lens API.
    Returns:
      {
        "font_name": str|None,
        "confidence": float,
        "alternatives": list[str],
        "image_url": str|None,
        "error": str|None
      }
    """
    if not MIXFONT_API_KEY:
        return {
            "font_name": None,
            "confidence": 0.0,
            "alternatives": [],
            "image_url": image_url,
            "error": "MIXFONT_API_KEY not set",
        }

    path_obj = Path(image_path)
    if not path_obj.exists():
        return {
            "font_name": None,
            "confidence": 0.0,
            "alternatives": [],
            "image_url": image_url,
            "error": f"image not found: {image_path}",
        }

    resolved_url = image_url or _build_public_image_url(str(path_obj))
    if not resolved_url:
        return {
            "font_name": None,
            "confidence": 0.0,
            "alternatives": [],
            "image_url": None,
            "error": "no public image URL (set MIXFONT_IMAGE_BASE_URL or MIXFONT_IMAGE_URL_TEMPLATE)",
        }

    payload = json.dumps({"image_url": resolved_url}).encode("utf-8")
    request = urllib.request.Request(
        MIXFONT_API_URL,
        data=payload,
        method="POST",
        headers={
            "Content-Type": "application/json",
            "x-api-key": MIXFONT_API_KEY,
        },
    )

    try:
        with urllib.request.urlopen(request, timeout=MIXFONT_TIMEOUT_SEC) as response:
            body = response.read().decode("utf-8", errors="replace")
        data = json.loads(body)
    except Exception as e:
        return {
            "font_name": None,
            "confidence": 0.0,
            "alternatives": [],
            "image_url": resolved_url,
            "error": str(e),
        }

    matches = data.get("font_matches", [])
    if not isinstance(matches, list):
        matches = []

    top = matches[0] if matches else {}
    if not isinstance(top, dict):
        top = {}

    font_name = str(top.get("name") or "").strip() or None
    try:
        confidence = float(top.get("confidence", 0.0))
    except Exception:
        confidence = 0.0

    alternatives = []
    for match in matches[1:4]:
        if isinstance(match, dict):
            name = str(match.get("name") or "").strip()
            if name:
                alternatives.append(name)

    if font_name:
        registry = _load_registry()
        entry = registry.get(font_name, _ensure_registry_entry(font_name))
        entry["identified_count"] = int(entry.get("identified_count", 0)) + 1
        entry["last_identified_at"] = str(data.get("datetime") or "")
        registry[font_name] = entry
        _save_registry(registry)

    return {
        "font_name": font_name,
        "confidence": round(confidence, 4),
        "alternatives": alternatives,
        "image_url": resolved_url,
        "error": None,
    }


def register_font_model(
    font_name: str,
    model_name: Optional[str] = None,
    preferred_engine: str = "tesseract",
    status: str = "available",
) -> dict:
    """
    Manually register a font->tesseract model mapping.
    Use this for bespoke fonts where only a specimen-based traineddata exists.
    """
    if not font_name:
        return {"ok": False, "error": "font_name required"}

    model = (model_name or _slug_font_name(font_name) or "eng").strip()
    registry = _load_registry()
    entry = registry.get(font_name, _ensure_registry_entry(font_name))
    entry["tesseract_model"] = model
    entry["preferred_engine"] = preferred_engine if preferred_engine in {"surya", "tesseract"} else "tesseract"
    entry["status"] = status
    registry[font_name] = entry
    _save_registry(registry)
    return {"ok": True, "font_name": font_name, "tesseract_model": model}


# Static font name → model name mapping.
# Handles cases where the slug doesn't match the traineddata filename.
FONT_MODEL_MAP: dict[str, str] = {
    "gill sans infant":  "gill_sans_infant",
    "gill sans":         "gill_sans_infant",  # MixFont may return either name
    "gill sans mt":      "gill_sans_infant",
    "gillsans":          "gill_sans_infant",
}


def resolve_tesseract_lang(font_name: Optional[str]) -> str:
    """
    Return best tesseract language/model to use for this font.
    Defaults to "eng" if no custom traineddata is available.
    """
    if not font_name:
        return "eng"

    # Check static map first
    mapped = FONT_MODEL_MAP.get(font_name.lower().strip())
    if mapped:
        tess_dirs = [TESSDATA_DIR, REPO_TESSDATA_DIR]
        if any((d / f"{mapped}.traineddata").exists() for d in tess_dirs):
            return mapped

    registry = _load_registry()
    entry = registry.get(font_name, {})
    candidates = []

    explicit_model = str(entry.get("tesseract_model") or "").strip()
    if explicit_model:
        candidates.append(explicit_model)

    slug = _slug_font_name(font_name)
    if slug:
        candidates.append(slug)
        candidates.append(slug.replace("_", ""))

    # unique keep order
    seen = set()
    deduped = []
    for c in candidates:
        if c and c not in seen:
            deduped.append(c)
            seen.add(c)

    tess_dirs = [TESSDATA_DIR, REPO_TESSDATA_DIR]
    for model_name in deduped:
        if any((tess_dir / f"{model_name}.traineddata").exists() for tess_dir in tess_dirs):
            if entry:
                entry["status"] = "available"
                entry["preferred_engine"] = "tesseract"
                entry["tesseract_model"] = model_name
                registry[font_name] = entry
                _save_registry(registry)
            return model_name

    # Common misplacement: .traineddata.otf (font file, not Tesseract model)
    otf_like_candidates = []
    for model_name in deduped:
        otf_like_candidates.extend(
            [
                REPO_TESSDATA_DIR / f"{model_name}.traineddata.otf",
                TESSDATA_DIR / f"{model_name}.traineddata.otf",
                REPO_TESSDATA_DIR / f"{model_name}.otf",
                TESSDATA_DIR / f"{model_name}.otf",
            ]
        )
    if any(path.exists() for path in otf_like_candidates):
        if entry:
            entry["status"] = "font_file_only"
            entry["preferred_engine"] = "surya"
            registry[font_name] = entry
            _save_registry(registry)

    # If a specimen exists but no model compiled yet, make that explicit in registry.
    if (FONTS_DIR / "gruffalo_specimen.png").exists():
        if entry:
            entry["status"] = entry.get("status") or "specimen_only"
            if entry["status"] == "unknown":
                entry["status"] = "specimen_only"
            registry[font_name] = entry
            _save_registry(registry)

    return "eng"


def resolve_tessdata_dir(model_name: Optional[str]) -> Optional[str]:
    if not model_name:
        return None
    model = str(model_name).strip()
    if not model:
        return None
    for tess_dir in (TESSDATA_DIR, REPO_TESSDATA_DIR):
        if (tess_dir / f"{model}.traineddata").exists():
            return str(tess_dir)
    return None


def update_font_engine_stats(font_name: Optional[str], surya_conf: float, tess_conf: float) -> None:
    if not font_name:
        return
    registry = _load_registry()
    entry = registry.get(font_name, _ensure_registry_entry(font_name))
    alpha = 0.3

    prev_surya = entry.get("avg_confidence_surya")
    prev_tess = entry.get("avg_confidence_tess")
    try:
        prev_surya_val = float(prev_surya) if prev_surya is not None else None
    except Exception:
        prev_surya_val = None
    try:
        prev_tess_val = float(prev_tess) if prev_tess is not None else None
    except Exception:
        prev_tess_val = None

    surya_new = float(surya_conf or 0.0)
    tess_new = float(tess_conf or 0.0)

    surya_avg = surya_new if prev_surya_val is None else round(alpha * surya_new + (1 - alpha) * prev_surya_val, 4)
    tess_avg = tess_new if prev_tess_val is None else round(alpha * tess_new + (1 - alpha) * prev_tess_val, 4)

    entry["avg_confidence_surya"] = surya_avg
    entry["avg_confidence_tess"] = tess_avg
    entry["preferred_engine"] = "tesseract" if tess_avg > surya_avg else "surya"
    if entry.get("status") == "unknown":
        entry["status"] = "available"

    registry[font_name] = entry
    _save_registry(registry)


def font_correction_map_path(font_name: Optional[str]) -> Path:
    """
    Font-specific punctuation learning map path.
    """
    if not font_name:
        return CALIBRATION_DIR / "punct_correction_map.json"
    safe = _slug_font_name(font_name) or "unknown"
    return CALIBRATION_DIR / f"punct_correction_map_{safe}.json"


def font_registry_summary() -> list[dict]:
    registry = _load_registry()
    rows = []
    for name, row in sorted(
        registry.items(),
        key=lambda x: int(x[1].get("identified_count", 0)),
        reverse=True,
    ):
        rows.append(
            {
                "font_name": name,
                "pages_seen": int(row.get("identified_count", 0)),
                "status": str(row.get("status", "unknown")),
                "preferred_engine": str(row.get("preferred_engine", "surya")),
                "tesseract_model": row.get("tesseract_model"),
                "surya_avg": row.get("avg_confidence_surya"),
                "tess_avg": row.get("avg_confidence_tess"),
            }
        )
    return rows