import numpy as np
import pytest
from PIL import Image

from pixtext_ocr import WordOCR


def quad(x0, y0, x1, y1):
    return np.float32([[x0, y0], [x1, y0], [x1, y1], [x0, y1]])


def routed_detector(parent, glyph_texts):
    engine = WordOCR.__new__(WordOCR)
    word = quad(0, 5, 150, 45)
    chars = [quad(x, 10, x + width, 40) for x, width in [(5, 19), (35, 24), (73, 18), (110, 30)]]
    engine.detect_quads = lambda *args: ([word], chars)
    calls = []

    def recognize(crops):
        if not crops:
            return []
        calls.append(len(crops))
        return [(parent, 0.95)] if len(calls) == 1 else [(s, 0.9) for s in glyph_texts]

    engine.recognize = recognize
    return engine, chars, calls


def test_chinese_uses_detected_glyph_boxes_and_independent_recognition():
    engine, chars, calls = routed_detector('新年快乐', ['新', '年', '怪', '乐'])
    rows = engine.detect(Image.new('RGB', (160, 60)), 0.5, 0.2)
    assert calls == [1, 4]
    assert [r['text'] for r in rows] == ['新', '年', '怪', '乐']
    assert [r['bbox'] for r in rows] == [[*q.min(axis=0), *q.max(axis=0)] for q in chars]
    assert all(r['level'] == 'character' and not r['enabled'] for r in rows)


def test_english_remains_a_whole_word():
    engine, _, calls = routed_detector('crossbody', [])
    rows = engine.detect(Image.new('RGB', (160, 60)), 0.5, 0.2)
    assert calls == [1]
    assert len(rows) == 1 and rows[0]['text'] == 'crossbody' and rows[0]['level'] == 'word'


def test_unreadable_or_ambiguous_glyphs_remain_editable_without_phrase_labels():
    engine, _, _ = routed_detector('新年快乐', ['', '新年', '快', '乐'])
    rows = engine.detect(Image.new('RGB', (160, 60)), 0.5, 0.2)
    assert len(rows) == 4
    assert [r['text'] for r in rows] == ['', '', '快', '乐']
    assert rows[0]['confidence'] == rows[1]['confidence'] == 0


def test_mixed_text_groups_latin_letters_but_keeps_han_individual():
    engine, _, _ = routed_detector('AB中文', ['A', 'B', '中', '文'])
    rows = engine.detect(Image.new('RGB', (160, 60)), 0.5, 0.2)
    assert [(r['text'], r['level']) for r in rows] == [('AB', 'word'), ('中', 'character'), ('文', 'character')]


def test_pure_chinese_does_not_merge_glyphs_misrecognized_as_latin():
    engine, _, _ = routed_detector('新年快乐', ['新', 'L', 'S', '乐'])
    rows = engine.detect(Image.new('RGB', (160, 60)), 0.5, 0.2)
    assert len(rows) == 4 and all(r['level'] == 'character' for r in rows)


def test_craft_affinity_joins_words_but_does_not_join_character_boxes():
    pytest.importorskip('easyocr')
    import torch

    scores = torch.zeros((1, 64, 128, 2))
    for x in [24, 55, 90]:
        scores[0, 24:40, x - 6:x + 6, 0] = 0.95
    scores[0, 29:35, 24:91, 1] = 0.9
    engine = WordOCR.__new__(WordOCR)
    engine.detector = lambda tensor: (scores, None)
    words, chars = engine.detect_quads(np.zeros((80, 160, 3), dtype=np.uint8), 0.5, 0.2)
    assert len(words) == 1 and len(chars) == 3
    centers = sorted(float(q.mean(axis=0)[0]) for q in chars)
    assert np.allclose(centers, np.array([24, 55, 90]) / 1.5 * 2, atol=3)