FeNigma/src/ironnest_assist/ocr.py

"""Text OCR pipeline: screenshot -> cleaned-up text -> parsed board info.

Only the text sub-pipeline is implemented. There's no image/icon
recognition sub-pipeline yet (spotting markers, ship icons, etc.) — that's
a separate future pipeline, out of scope here.

Preprocessing matters more than the regexes: the typewriter photo has an
uneven vignette (in-game light falloff) that a single global threshold
can't handle — it either loses faint corners or blobs-out dark ones. We
flatten that by dividing by a heavily blurred copy of itself (crude local
background normalization) before thresholding, which recovers text in the
darkened areas reliably.

Parsing is deliberately fuzzy: OCR on an in-game screenshot will misread
the odd character (keywords slightly garbled, '0'/'O', '1'/'I'/'l',
'5'/'S', '8'/'B' confused, stray specks turning ':' into ';' or '.'). We
match keywords by similarity rather than exact spelling, and normalize
digit-shaped letters before parsing numbers, so a couple of misrecognized
characters don't drop an otherwise-good line.
"""

from __future__ import annotations

import difflib
import re
from dataclasses import dataclass, field

import numpy as np
import pytesseract
from PIL import Image, ImageFilter

from .models import Clue, Coord, TargetType

# --- preprocessing -----------------------------------------------------------

_BLUR_RADIUS = 35
_THRESHOLD = 200


def preprocess(image: Image.Image) -> Image.Image:
    gray = image.convert("L")
    arr = np.asarray(gray, dtype=np.float32)
    background = np.asarray(gray.filter(ImageFilter.GaussianBlur(_BLUR_RADIUS)), dtype=np.float32)
    normalized = np.clip(arr / (background + 1e-3) * 255.0, 0, 255).astype(np.uint8)
    normalized_img = Image.fromarray(normalized)
    return normalized_img.point(lambda p: 255 if p > _THRESHOLD else 0)


def ocr_text(image: Image.Image) -> str:
    return pytesseract.image_to_string(preprocess(image), config="--psm 6")


# --- fuzzy keyword matching ---------------------------------------------------

NEST_KEYWORD = "IRON NEST"
SPOTTER_KEYWORD = "SPOTTER"
_FUZZY_THRESHOLD = 0.65


def _fuzzy_locate(line: str, keyword: str, threshold: float = _FUZZY_THRESHOLD) -> tuple[int, int] | None:
    """Span of the best approximate match of `keyword` in `line`, or None."""
    line_u = line.upper()
    kw = keyword.upper()
    n = len(kw)
    best_ratio = 0.0
    best_span = None
    for size in range(max(n - 2, 1), n + 3):
        for start in range(0, max(len(line_u) - size, 0) + 1):
            end = start + size
            ratio = difflib.SequenceMatcher(None, line_u[start:end], kw).ratio()
            if ratio > best_ratio:
                best_ratio = ratio
                best_span = (start, end)
    return best_span if best_ratio >= threshold else None


def _fuzzy_contains(line: str, keyword: str, threshold: float = _FUZZY_THRESHOLD) -> bool:
    """True if some run of characters in `line` approximately matches `keyword`."""
    return _fuzzy_locate(line, keyword, threshold) is not None


# --- coordinate / id extraction, tolerant of digit<->letter OCR mixups --------

_DIGIT_CLASS = r"[0-9OoIiLlSsBbZzGg]"
_DIGIT_FIX = str.maketrans({
    "O": "0", "o": "0",
    "I": "1", "i": "1", "L": "1", "l": "1",
    "S": "5", "s": "5",
    "B": "8", "b": "8",
    "Z": "2", "z": "2",
    "G": "6", "g": "6",
})
_SEP = r"\s*[-–—]\s*"
_COORD_RE = re.compile(
    rf"([A-T])\s*({_DIGIT_CLASS}{{1,2}})\s+({_DIGIT_CLASS})\s*[:;.,]\s*({_DIGIT_CLASS})"
)
# No literal '#' required — it's just as OCR-corruptible as anything else
# (missing entirely, or misread as e.g. 'H'). We instead anchor to *where*
# the fuzzy keyword match ended and take the first run of digit-shaped
# characters after that, skipping over whatever separator survived.
_ID_RE = re.compile(rf"({_DIGIT_CLASS}+)")

# Not seen in a real screenshot yet, so no extraction for these — add a
# keyword + extractor here (plus a field below and a case in parse_text)
# once we know the format:
#   - reference points
#   - targets


def _fix_digits(s: str) -> str:
    return s.translate(_DIGIT_FIX)


def _fix_id_digits(raw: str) -> str:
    """Like _fix_digits, but for id-length runs specifically: also
    collapses a 2-character run where one char is a genuine digit and the
    other a look-alike letter that maps to the *same* digit — e.g. '1l' or
    'S5' both fix to '11'/'55', but a real 2-digit id wouldn't plausibly
    render as one numeral plus one letter of the identical value; that
    shape is the signature of OCR ghosting a single thin glyph twice
    (seen repeatedly: 'AmmoCache#l1' -> #1, 'HostileTank#S5' -> #5).
    A genuine two-digit id (both chars already real digits, e.g. '11')
    is left alone."""
    fixed = _fix_digits(raw)
    if len(raw) == 2 and len(set(fixed)) == 1:
        has_digit = any(c.isdigit() for c in raw)
        has_letter = any(c.isalpha() for c in raw)
        if has_digit and has_letter:
            return fixed[0]
    return fixed


# A target can also be spotted with an absolute grid ref directly
# ("Target#10 Spotted. Grid Q3 9:0") instead of/alongside bearing/distance
# clues — same coordinate shape as _COORD_RE, just anchored after "Grid".
_GRID_COORD_RE = re.compile(
    rf"Grid\s+([A-T])\s*({_DIGIT_CLASS}{{1,2}})\s+({_DIGIT_CLASS})\s*[:;.,]\s*({_DIGIT_CLASS})",
    re.IGNORECASE,
)


def _extract_grid_coord(text: str) -> Coord | None:
    m = _GRID_COORD_RE.search(text)
    if not m:
        return None
    letter, y, x, yy = m.groups()
    try:
        return Coord(X=letter.upper(), Y=int(_fix_digits(y)), x=int(_fix_digits(x)), y=int(_fix_digits(yy)))
    except ValueError:
        return None


def _extract_coord(line: str) -> Coord | None:
    m = _COORD_RE.search(line)
    if not m:
        return None
    letter, y, x, yy = m.groups()
    try:
        return Coord(
            X=letter.upper(),
            Y=int(_fix_digits(y)),
            x=int(_fix_digits(x)),
            y=int(_fix_digits(yy)),
        )
    except ValueError:
        return None  # out-of-range numbers -> not actually a coordinate


def _extract_leading_id(remainder: str) -> int | None:
    """First digit-shaped run in `remainder` (text after the keyword match)."""
    m = _ID_RE.search(remainder)
    if not m:
        return None
    try:
        return int(_fix_id_digits(m.group(1)))
    except ValueError:
        return None


# --- "field intelligence" blocks: relative Bearing/Distance descriptions ------
#
#   Target#5  Spotted. 088, 12.10km from Spotter#1
#   Reference Point Alpha:
#       Bearing 094 from Spotter#1
#       Distance 13.26km from Spotter#2
#   .
#   AmmoCache#3:
#       Bearing 217 & Distance 10.48km from AmmoCache#2
#
# Each named entity is a block of one or more clue lines, terminated by a
# blank line or a lone '.'. Degree signs, colons, and the 'km' unit are all
# treated as optional/lossy — OCR drops them unpredictably.

_RP_HEADER_RE = re.compile(r"Reference\s+Point\s+([A-Za-z][\w-]*)\s*:?", re.IGNORECASE)
# '#' isn't required literally — same reasoning as the spotter-id fix: OCR
# drops it or renders it as noise (seen: '€'). Up to 2 junk characters
# between the type word and its digits is enough slack without risking a
# false match elsewhere.
#   ^ the junk-class run is REQUIRED (1-2 chars, not 0-2): the digit class
#   deliberately overlaps the alphabet (g/s/i/l/o/... look like digits), so
#   with a 0-width separator allowed, "Bearing" backtracks into itself —
#   word="Bearin", "digit"=its own trailing 'g' — and falsely matches as a
#   header. Requiring real punctuation between word and digits (true of
#   every observed header: '#', a misread substitute, ...) rules that out,
#   and also stops a bare "Word 094" clue line (space only) from matching.
_NAMED_HEADER_RE = re.compile(rf"^([A-Za-z]+)[^A-Za-z0-9\s]{{1,2}}({_DIGIT_CLASS}+)\s*:?\s*(.*)$")
_BLOCK_END_RE = re.compile(r"^[.\s]{1,6}$")
_LEADING_NOISE_RE = re.compile(r"^[^A-Za-z]{1,3}(?=[A-Za-z])")

# Reference capture is (\S+), not (.+): references are always a single
# token with no spaces, and being non-greedy this way is what lets
# finditer() find more than one clue per line/block — "Bearing 118 from
# Spotter#2 & Bearing 125 from Spotter#1" needs two separate matches, and
# a greedy (.+) would let the first one swallow the rest of the string.
#
# _GAP sits right before "from": plain whitespace normally, but an OCR
# line-wrap can drop a stray junk token right at the break ("Bearing 125°"
# / "P; from Spotter#1") — and that junk can itself contain a letter (the
# 'P' above), so this isn't just non-alnum noise like the header-bullet
# case; tolerate any single short token, not just punctuation.
_GAP = r"[\s]*(?:\S{1,3}\s*)?"
_CLUE_COMBINED_RE = re.compile(
    rf"Bearing\s*(\d{{1,3}})\s*°?\s*&\s*Distance\s*([\d.]+)\s*k?m?{_GAP}from\s+(\S+)", re.IGNORECASE
)
_CLUE_INLINE_RE = re.compile(
    rf"(\d{{1,3}})\s*°?\s*,\s*([\d.]+)\s*k?m{_GAP}from\s+(\S+)", re.IGNORECASE
)
_CLUE_BEARING_RE = re.compile(rf"Bearing\s*(\d{{1,3}})\s*°?{_GAP}from\s+(\S+)", re.IGNORECASE)
_CLUE_DISTANCE_RE = re.compile(rf"Distance\s*([\d.]+)\s*k?m?{_GAP}from\s+(\S+)", re.IGNORECASE)

_TYPE_BY_SHORT = {t.short: t for t in TargetType}
# The game's typewriter has used "AmmoCache" for what's now modeled as
# SupplyCache — treat it as the same type rather than dropping the target.
_TYPE_WORD_ALIASES = {"AmmoCache": "SupplyCache"}


_REF_NAMED_RE = re.compile(rf"^([A-Za-z]+)[^A-Za-z0-9\s]{{1,2}}({_DIGIT_CLASS}+)")


def _clean_reference(raw: str) -> str:
    """Leading name-shaped token, dropping trailing OCR noise (stray dots,
    double spaces, etc.) — reference names never contain spaces. Named
    references ('Spotter#1', 'AmmoCache#2') get their digit part fixed up
    and their separator normalized to '#'; plain word references ('Alpha')
    are left untouched — don't run digit-fixing over them or real letters
    like the 'l' in 'Alpha' get corrupted into '1'."""
    token = re.match(r"\S+", raw.strip())
    token = token.group(0) if token else raw.strip()

    named = _REF_NAMED_RE.match(token)
    if named:
        word, num = named.groups()
        word = _TYPE_WORD_ALIASES.get(word, word)
        return f"{word}#{_fix_id_digits(num)}"

    return re.sub(r"[^\w-]+$", "", token)  # trim trailing punctuation off a plain name


# Priority order matters: try the "both bearing and distance" shape before
# the single-value shapes, so e.g. "Bearing 217 & Distance 10.48km from X"
# becomes one clue, not a spurious extra bearing-only one from the same
# text. `\s` already matches a literal newline, so this bridges an OCR
# line-wrap ("...Bearing 125°" / "from Spotter#1" split across two lines)
# without needing the caller to rejoin anything.
_CLUE_PATTERNS = (
    (_CLUE_COMBINED_RE, lambda m: (float(m.group(1)), float(m.group(2)), m.group(3))),
    (_CLUE_INLINE_RE, lambda m: (float(m.group(1)), float(m.group(2)), m.group(3))),
    (_CLUE_BEARING_RE, lambda m: (float(m.group(1)), None, m.group(2))),
    (_CLUE_DISTANCE_RE, lambda m: (None, float(m.group(1)), m.group(2))),
)


def _parse_all_clues(text: str) -> list[Clue]:
    """Every Bearing/Distance clue found anywhere in `text` — a block can
    have several (one per clue line, or more than one on a single line
    joined with '&')."""
    clues: list[Clue] = []
    consumed: list[tuple[int, int]] = []

    def overlaps(span: tuple[int, int]) -> bool:
        return any(span[0] < e and span[1] > s for s, e in consumed)

    for pattern, extract in _CLUE_PATTERNS:
        for m in pattern.finditer(text):
            span = m.span()
            if overlaps(span):
                continue
            consumed.append(span)
            bearing, distance, ref = extract(m)
            clues.append(Clue(reference=_clean_reference(ref), bearing_deg=bearing, distance_km=distance))
    return clues


def parse_clues_from_text(text: str) -> list[Clue]:
    """Parse every Bearing/Distance clue found in free-form text — used
    for manually-typed descriptions in the coord dialog, sharing the
    exact same clue grammar as the OCR'd intel blocks."""
    return _parse_all_clues(text)


def _resolve_target_type(type_word: str) -> TargetType | None:
    """Exact match on the type word (after aliasing), falling back to
    fuzzy (OCR can garble the type word itself, e.g. 'AmmoCoche')."""
    type_word = _TYPE_WORD_ALIASES.get(type_word, type_word)
    if type_word in _TYPE_BY_SHORT:
        return _TYPE_BY_SHORT[type_word]
    best, best_ratio = None, 0.0
    for short, target_type in _TYPE_BY_SHORT.items():
        ratio = difflib.SequenceMatcher(None, type_word.upper(), short.upper()).ratio()
        if ratio > best_ratio:
            best, best_ratio = target_type, ratio
    return best if best_ratio >= _FUZZY_THRESHOLD else None


def parse_intel_blocks(text: str) -> list[dict]:
    """Parse 'field intelligence' blocks into a list of dicts with keys
    kind ('rp' | 'named'), name, type_word, id, raw, clues, coord. Entries
    with neither a clue nor a grid coord are dropped (nothing to store).

    Clues are extracted once per block, from the whole joined block text,
    at flush time — not accumulated line-by-line while scanning. That's
    what lets a clue split across an OCR line-wrap ("...Bearing 125°" /
    "from Spotter#1" on separate lines) or two clues on one line
    ("Bearing X from A & Bearing Y from B") both resolve correctly. A
    target can also carry an absolute grid ref directly ("Grid Q3 9:0")
    instead of/alongside clues."""
    entries: list[dict] = []
    current: dict | None = None

    def flush():
        nonlocal current
        if current is not None:
            joined = "\n".join(current["raw"])
            current["clues"] = _parse_all_clues(joined)
            current["coord"] = _extract_grid_coord(joined)
            if current["clues"] or current["coord"] is not None:
                current["raw"] = joined
                entries.append(current)
        current = None

    for raw_line in text.splitlines():
        line = raw_line.strip()
        if not line:
            continue  # blank lines are just visual spacing here, not a section boundary
        if _BLOCK_END_RE.match(line):
            flush()
            continue

        # The '.' block-separator bullet often survives OCR as 1-3 stray
        # junk characters glued onto the *next* line ('i AmmoCache#1:',
        # '* AmmoCache#2:') instead of its own line, which would otherwise
        # defeat the column-0-anchored header regexes below. Try the line
        # as-is first, and only if that fails, retry with its first
        # whitespace-delimited token stripped (covers symbol junk *and*
        # a misread bullet that happened to OCR as a stray letter) — but
        # only when that leading token is bullet-length (<=3 chars), else
        # a genuine wrapped clue continuation like "from Spotter#2" gets
        # its "from" stripped and "Spotter#2" misread as a new header.
        first_token_m = re.match(r"\S+", line)
        strip_first_token = (
            re.sub(r"^\S+\s+", "", line, count=1)
            if first_token_m and len(first_token_m.group(0)) <= 3
            else line
        )
        candidates = (line, _LEADING_NOISE_RE.sub("", line), strip_first_token)

        rp_m = next((m for c in candidates if (m := _RP_HEADER_RE.search(c))), None)
        if rp_m:
            flush()
            current = {"kind": "rp", "name": rp_m.group(1), "type_word": None, "id": None,
                       "raw": [line], "clues": []}
            continue

        named_m = next((m for c in candidates if (m := _NAMED_HEADER_RE.match(c))), None)
        if named_m:
            flush()
            type_word, num, inline = named_m.groups()
            num = _fix_id_digits(num)
            current = {"kind": "named", "name": f"{type_word}#{num}", "type_word": type_word,
                       "id": num, "raw": [line], "clues": []}
            continue

        if current is not None:
            current["raw"].append(line)

    flush()
    return entries


# Destruction reports are standalone one-liners, not tied to a block:
#   "SupplyCache#2 Destroyed. Additional Requisition Granted."
#   "Direct Hit! HostileTank#3 Destroyed."
# Just need "<Type>#<id>" immediately followed by "Destroyed" — the
# "Direct Hit!" prefix (or its absence) doesn't matter, search() finds
# the name+Destroyed pair anywhere in the line either way.
_DESTROYED_RE = re.compile(
    rf"([A-Za-z]+)[^A-Za-z0-9\s]{{1,2}}({_DIGIT_CLASS}+)\s*Destroyed", re.IGNORECASE
)


def parse_destroyed(text: str) -> set[tuple[TargetType, str]]:
    destroyed = set()
    for m in _DESTROYED_RE.finditer(text):
        type_word, num = m.groups()
        target_type = _resolve_target_type(type_word)
        if target_type is None:
            continue
        destroyed.add((target_type, _fix_id_digits(num)))
    return destroyed


@dataclass
class ParsedInfo:
    nest_coord: Coord | None = None
    spotters: dict[int, Coord] = field(default_factory=dict)
    # name -> (raw description, clues, absolute coord if given directly)
    reference_points: dict[str, tuple[str, list[Clue], Coord | None]] = field(default_factory=dict)
    # (type, id) -> (raw description, clues, absolute coord if given directly)
    targets: dict[tuple[TargetType, str], tuple[str, list[Clue], Coord | None]] = field(default_factory=dict)
    # (type, id) of targets reported destroyed
    destroyed: set[tuple[TargetType, str]] = field(default_factory=set)

    def is_empty(self) -> bool:
        return not (self.nest_coord or self.spotters or self.reference_points
                    or self.targets or self.destroyed)


def parse_text(text: str) -> ParsedInfo:
    info = ParsedInfo()

    for raw_line in text.splitlines():
        line = raw_line.strip()
        if not line:
            continue

        if info.nest_coord is None and _fuzzy_contains(line, NEST_KEYWORD):
            coord = _extract_coord(line)
            if coord is not None:
                info.nest_coord = coord
                continue

        spotter_span = _fuzzy_locate(line, SPOTTER_KEYWORD)
        if spotter_span is not None:
            spotter_id = _extract_leading_id(line[spotter_span[1]:])
            coord = _extract_coord(line)
            if spotter_id is not None and coord is not None:
                info.spotters[spotter_id] = coord

    for entry in parse_intel_blocks(text):
        if entry["kind"] == "rp":
            info.reference_points[entry["name"]] = (entry["raw"], entry["clues"], entry["coord"])
            continue
        target_type = _resolve_target_type(entry["type_word"])
        if target_type is None:
            continue
        info.targets[(target_type, entry["id"])] = (entry["raw"], entry["clues"], entry["coord"])

    info.destroyed = parse_destroyed(text)
    return info


def run(image: Image.Image) -> ParsedInfo:
    return parse_text(ocr_text(image))