"""
OCR Error Corrections Module

Provides comprehensive OCR error correction tables and correction functions.
Based on common OCR recognition errors in Swedish invoice documents.
"""

import re
from typing import Optional
from dataclasses import dataclass


@dataclass
class CorrectionResult:
    """Result of an OCR correction operation."""
    original: str
    corrected: str
    corrections_applied: list[tuple[int, str, str]]  # (position, from_char, to_char)
    confidence: float  # How confident we are in the correction


class OCRCorrections:
    """
    Comprehensive OCR error correction utilities.

    Provides:
    - Character-level corrections for digits
    - Word-level corrections for common Swedish terms
    - Context-aware corrections
    - Multiple correction strategies
    """

    # =========================================================================
    # Character-level OCR errors (digit fields)
    # =========================================================================

    # Characters commonly misread as digits
    CHAR_TO_DIGIT = {
        # Letters that look like digits
        'O': '0', 'o': '0',      # O -> 0
        'Q': '0',                 # Q -> 0 (less common)
        'D': '0',                 # D -> 0 (in some fonts)

        'l': '1', 'I': '1',      # l/I -> 1
        'i': '1',                 # i without dot -> 1
        '|': '1',                 # pipe -> 1
        '!': '1',                 # exclamation -> 1

        'Z': '2', 'z': '2',      # Z -> 2

        'E': '3',                 # E -> 3 (rare)

        'A': '4', 'h': '4',      # A/h -> 4 (in some fonts)

        'S': '5', 's': '5',      # S -> 5

        'G': '6', 'b': '6',      # G/b -> 6

        'T': '7', 't': '7',      # T -> 7 (rare)

        'B': '8',                 # B -> 8

        'g': '9', 'q': '9',      # g/q -> 9
    }

    # Digits commonly misread as other characters
    DIGIT_TO_CHAR = {
        '0': ['O', 'o', 'D', 'Q'],
        '1': ['l', 'I', 'i', '|', '!'],
        '2': ['Z', 'z'],
        '3': ['E'],
        '4': ['A', 'h'],
        '5': ['S', 's'],
        '6': ['G', 'b'],
        '7': ['T', 't'],
        '8': ['B'],
        '9': ['g', 'q'],
    }

    # Bidirectional confusion pairs (either direction is possible)
    CONFUSION_PAIRS = [
        ('0', 'O'), ('0', 'o'), ('0', 'D'),
        ('1', 'l'), ('1', 'I'), ('1', '|'),
        ('2', 'Z'), ('2', 'z'),
        ('5', 'S'), ('5', 's'),
        ('6', 'G'), ('6', 'b'),
        ('8', 'B'),
        ('9', 'g'), ('9', 'q'),
    ]

    # =========================================================================
    # Word-level OCR errors (Swedish invoice terms)
    # =========================================================================

    # Common Swedish invoice terms and their OCR misreadings
    SWEDISH_TERM_CORRECTIONS = {
        # Faktura (Invoice)
        'faktura': ['Faktura', 'FAKTURA', 'faktúra', 'faKtura'],
        'fakturanummer': ['Fakturanummer', 'FAKTURANUMMER', 'fakturanr', 'fakt.nr'],
        'fakturadatum': ['Fakturadatum', 'FAKTURADATUM', 'fakt.datum'],

        # Belopp (Amount)
        'belopp': ['Belopp', 'BELOPP', 'be1opp', 'bel0pp'],
        'summa': ['Summa', 'SUMMA', '5umma'],
        'total': ['Total', 'TOTAL', 'tota1', 't0tal'],
        'moms': ['Moms', 'MOMS', 'm0ms'],

        # Dates
        'förfallodatum': ['Förfallodatum', 'FÖRFALLODATUM', 'förfa11odatum'],
        'datum': ['Datum', 'DATUM', 'dátum'],

        # Payment
        'bankgiro': ['Bankgiro', 'BANKGIRO', 'BG', 'bg', 'bank giro'],
        'plusgiro': ['Plusgiro', 'PLUSGIRO', 'PG', 'pg', 'plus giro'],
        'postgiro': ['Postgiro', 'POSTGIRO'],
        'ocr': ['OCR', 'ocr', '0CR', 'OcR'],

        # Organization
        'organisationsnummer': ['Organisationsnummer', 'ORGANISATIONSNUMMER', 'org.nr', 'orgnr'],
        'kundnummer': ['Kundnummer', 'KUNDNUMMER', 'kund nr', 'kundnr'],

        # Currency
        'kronor': ['Kronor', 'KRONOR', 'kr', 'KR', 'SEK', 'sek'],
        'öre': ['Öre', 'ÖRE', 'ore', 'ORE'],
    }

    # =========================================================================
    # Context patterns
    # =========================================================================

    # Patterns that indicate the following/preceding text is a specific field
    CONTEXT_INDICATORS = {
        'invoice_number': [
            r'faktura\s*(?:nr|nummer)?[:\s]*',
            r'invoice\s*(?:no|number)?[:\s]*',
            r'fakt\.?\s*nr[:\s]*',
            r'inv[:\s]*#?',
        ],
        'amount': [
            r'(?:att\s+)?betala[:\s]*',
            r'total[t]?[:\s]*',
            r'summa[:\s]*',
            r'belopp[:\s]*',
            r'amount[:\s]*',
        ],
        'date': [
            r'datum[:\s]*',
            r'date[:\s]*',
            r'förfall(?:o)?datum[:\s]*',
            r'fakturadatum[:\s]*',
        ],
        'ocr': [
            r'ocr[:\s]*',
            r'referens[:\s]*',
            r'betalningsreferens[:\s]*',
        ],
        'bankgiro': [
            r'bankgiro[:\s]*',
            r'bg[:\s]*',
            r'bank\s*giro[:\s]*',
        ],
        'plusgiro': [
            r'plusgiro[:\s]*',
            r'pg[:\s]*',
            r'plus\s*giro[:\s]*',
            r'postgiro[:\s]*',
        ],
        'org_number': [
            r'org\.?\s*(?:nr|nummer)?[:\s]*',
            r'organisationsnummer[:\s]*',
            r'vat[:\s]*',
            r'moms(?:reg)?\.?\s*(?:nr|nummer)?[:\s]*',
        ],
    }

    # =========================================================================
    # Correction Methods
    # =========================================================================

    @classmethod
    def correct_digits(cls, text: str, aggressive: bool = False) -> CorrectionResult:
        """
        Apply digit corrections to text.

        Args:
            text: Input text
            aggressive: If True, correct all potential digit-like characters.
                       If False, only correct characters adjacent to digits.

        Returns:
            CorrectionResult with original, corrected text, and details.
        """
        corrections = []
        result = []

        for i, char in enumerate(text):
            if char.isdigit():
                result.append(char)
            elif char in cls.CHAR_TO_DIGIT:
                if aggressive:
                    # Always correct
                    corrected_char = cls.CHAR_TO_DIGIT[char]
                    corrections.append((i, char, corrected_char))
                    result.append(corrected_char)
                else:
                    # Only correct if adjacent to digit
                    prev_is_digit = i > 0 and (text[i-1].isdigit() or text[i-1] in cls.CHAR_TO_DIGIT)
                    next_is_digit = i < len(text) - 1 and (text[i+1].isdigit() or text[i+1] in cls.CHAR_TO_DIGIT)

                    if prev_is_digit or next_is_digit:
                        corrected_char = cls.CHAR_TO_DIGIT[char]
                        corrections.append((i, char, corrected_char))
                        result.append(corrected_char)
                    else:
                        result.append(char)
            else:
                result.append(char)

        corrected = ''.join(result)
        confidence = 1.0 - (len(corrections) * 0.05)  # Decrease confidence per correction

        return CorrectionResult(
            original=text,
            corrected=corrected,
            corrections_applied=corrections,
            confidence=max(0.5, confidence)
        )

    @classmethod
    def generate_digit_variants(cls, text: str) -> list[str]:
        """
        Generate all possible digit interpretations of a text.

        Useful for matching when we don't know which direction the OCR error went.
        """
        if not text:
            return [text]

        variants = {text}

        # For each character that could be confused
        for i, char in enumerate(text):
            new_variants = set()
            for existing in variants:
                # If it's a digit, add letter variants
                if char.isdigit() and char in cls.DIGIT_TO_CHAR:
                    for replacement in cls.DIGIT_TO_CHAR[char]:
                        new_variants.add(existing[:i] + replacement + existing[i+1:])

                # If it's a letter that looks like a digit, add digit variant
                if char in cls.CHAR_TO_DIGIT:
                    new_variants.add(existing[:i] + cls.CHAR_TO_DIGIT[char] + existing[i+1:])

            variants.update(new_variants)

            # Limit explosion - only keep reasonable number
            if len(variants) > 100:
                break

        return list(variants)

    @classmethod
    def correct_swedish_term(cls, text: str) -> str:
        """
        Correct common Swedish invoice terms that may have OCR errors.
        """
        text_lower = text.lower()

        for canonical, variants in cls.SWEDISH_TERM_CORRECTIONS.items():
            for variant in variants:
                if variant.lower() in text_lower:
                    # Replace with canonical form (preserving case of first letter)
                    pattern = re.compile(re.escape(variant), re.IGNORECASE)
                    if text[0].isupper():
                        replacement = canonical.capitalize()
                    else:
                        replacement = canonical
                    text = pattern.sub(replacement, text)

        return text

    @classmethod
    def extract_with_context(cls, text: str, field_type: str) -> Optional[str]:
        """
        Extract a field value using context indicators.

        Looks for patterns like "Fakturanr: 12345" and extracts "12345".
        """
        patterns = cls.CONTEXT_INDICATORS.get(field_type, [])

        for pattern in patterns:
            # Look for pattern followed by value
            full_pattern = pattern + r'([^\s,;]+)'
            match = re.search(full_pattern, text, re.IGNORECASE)
            if match:
                return match.group(1)

        return None

    @classmethod
    def is_likely_ocr_error(cls, char1: str, char2: str) -> bool:
        """
        Check if two characters are commonly confused in OCR.
        """
        pair = (char1, char2)
        reverse_pair = (char2, char1)

        for p in cls.CONFUSION_PAIRS:
            if pair == p or reverse_pair == p:
                return True

        return False

    @classmethod
    def count_potential_ocr_errors(cls, s1: str, s2: str) -> tuple[int, int]:
        """
        Count how many character differences between two strings
        are likely OCR errors vs other differences.

        Returns: (ocr_errors, other_errors)
        """
        if len(s1) != len(s2):
            return (0, abs(len(s1) - len(s2)))

        ocr_errors = 0
        other_errors = 0

        for c1, c2 in zip(s1, s2):
            if c1 != c2:
                if cls.is_likely_ocr_error(c1, c2):
                    ocr_errors += 1
                else:
                    other_errors += 1

        return (ocr_errors, other_errors)

    @classmethod
    def suggest_corrections(cls, text: str, expected_type: str = 'digit') -> list[tuple[str, float]]:
        """
        Suggest possible corrections for a text with confidence scores.

        Returns list of (corrected_text, confidence) tuples, sorted by confidence.
        """
        suggestions = []

        if expected_type == 'digit':
            # Apply digit corrections with different levels of aggressiveness
            mild = cls.correct_digits(text, aggressive=False)
            if mild.corrected != text:
                suggestions.append((mild.corrected, mild.confidence))

            aggressive = cls.correct_digits(text, aggressive=True)
            if aggressive.corrected != text and aggressive.corrected != mild.corrected:
                suggestions.append((aggressive.corrected, aggressive.confidence * 0.9))

            # Generate variants
            variants = cls.generate_digit_variants(text)
            for variant in variants[:10]:  # Limit to top 10
                if variant != text and variant not in [s[0] for s in suggestions]:
                    # Lower confidence for variants
                    suggestions.append((variant, 0.7))

        # Sort by confidence
        suggestions.sort(key=lambda x: x[1], reverse=True)
        return suggestions


# =========================================================================
# Convenience functions
# =========================================================================

def correct_ocr_digits(text: str, aggressive: bool = False) -> str:
    """Convenience function to correct OCR digit errors."""
    return OCRCorrections.correct_digits(text, aggressive).corrected


def generate_ocr_variants(text: str) -> list[str]:
    """Convenience function to generate OCR variants."""
    return OCRCorrections.generate_digit_variants(text)


def is_ocr_confusion(char1: str, char2: str) -> bool:
    """Convenience function to check if characters are OCR confusable."""
    return OCRCorrections.is_likely_ocr_error(char1, char2)
