import cv2
import easyocr
import torch
import gc

tesseract_rotation_config = "-l eng+hin --psm 6"

# Shared EasyOCR reader instance
_easyocr_reader = None

def get_easyocr_reader():
    """Unified access to the EasyOCR reader to save memory."""
    global _easyocr_reader
    if _easyocr_reader is None:
        _easyocr_reader = easyocr.Reader(['en'], gpu=False)
    return _easyocr_reader

import pandas as pd
import re
import pytesseract

def ocr_score2(image):
    """Calculate an OCR confidence score to determine best rotation."""
    data = pytesseract.image_to_data(image, config=tesseract_rotation_config, output_type=pytesseract.Output.DICT)

    words = [
        (text.strip(), int(conf))
        for text, conf in zip(data["text"], data["conf"])
        if conf != "-1" and text.strip()
    ]

    if not words:
        return 0

    total_conf = 0
    valid_word_count = 0

    for text, conf in words:
        if re.match(r'^[A-Za-z0-9]{3,}$', text):
            total_conf += conf
            valid_word_count += 1

    if valid_word_count == 0:
        return 0

    mean_conf = total_conf / valid_word_count
    return mean_conf * (valid_word_count ** 0.5)

def pan_ocr_score(image):
    """Calculate a score for PAN card orientation based on keywords and confidence."""
    # Apply a light blur to reduce noise for the orientation check
    proc = cv2.GaussianBlur(image, (3, 3), 0)
    
    # Use Hin+Eng as many PAN cards have bilingual headers
    config = "-l eng+hin --psm 6"
    data = pytesseract.image_to_data(proc, config=config, output_type=pytesseract.Output.DICT)
    
    text_content = " ".join([t.strip().upper() for t in data["text"] if t.strip()])
    
    # Strong PAN markers (Phrases)
    phrases = ["INCOME TAX", "TAX DEPARTMENT", "ACCOUNT NUMBER", "CARD", "INDIA", "GOVT"]
    kw_score = 0
    for ph in phrases:
        if ph in text_content:
            kw_score += 500 # Heavy weight for phrases
            
    # Single word keywords with boundary check
    keywords = ["DEPARTMENT", "PERMANENT", "ACCOUNT", "NUMBER", "INCOME"]
    for kw in keywords:
        if re.search(r'\b' + kw + r'\b', text_content):
            kw_score += 100
            
    # Confidence-based score
    words = [int(conf) for conf in data["conf"] if conf != "-1"]
    if not words:
        return kw_score
    
    avg_conf = sum(words) / len(words)
    return kw_score + (avg_conf * (len(words) ** 0.5))

def extract_aadhaar_easy_ocr(img, app_id, logger=None):
    """Fallback OCR using EasyOCR for Aadhaar number extraction."""
    if logger:
        logger.send_log_event(f"Application_id : {app_id}. In Easy OCR Aadhaar extract function.")
    
    reader = get_easyocr_reader()
    easy_ocr_result = reader.readtext(img, allowlist='0123456789', contrast_ths=0.05, adjust_contrast=0.7, paragraph=False)
    
    data = []
    for (bbox, text, prob) in easy_ocr_result:
        data.append({"bbox": bbox, "text": text, "confidence": prob})
    
    df = pd.DataFrame(data)
    filtered = df[(df['text'].str.fullmatch(r'\d{12}', na=False)) & (df['confidence'] > 0.9)]
    
    if not filtered.empty:
        row = filtered.iloc[0]
        # bbox is [[x0, y0], [x1, y1], [x2, y2], [x3, y3]]
        return row['text'], [row['bbox']]
    return None, None

def find_best_rotation(gray_img):
    """Find the best 90-degree rotation based on OCR confidence."""
    gray_img = cv2.GaussianBlur(gray_img, (5, 5), 0)
    thresh = cv2.threshold(gray_img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1]
    
    rotations = {
        0: thresh,
        90: cv2.rotate(thresh, cv2.ROTATE_90_CLOCKWISE),
        180: cv2.rotate(thresh, cv2.ROTATE_180),
        270: cv2.rotate(thresh, cv2.ROTATE_90_COUNTERCLOCKWISE)
    }

    best_angle = 0
    best_score = -1
    best_img = thresh

    for angle, img_rot in rotations.items():
        score = ocr_score2(img_rot)
        if score > best_score:
            best_score = score
            best_angle = angle
            best_img = img_rot

    return best_angle, best_img

def find_best_rotation_pan(gray_img):
    """
    Specialized rotation finder for PAN cards using keyword weights.
    Efficiently short-circuits if 0-degree orientation is detected with high confidence.
    """
    best_angle = 0
    best_score = -1
    best_img = gray_img

    # Check rotations sequentially to allow for early exit
    for angle in [0, 90, 180, 270]:
        if angle == 0:
            rotated = gray_img
        elif angle == 90:
            rotated = cv2.rotate(gray_img, cv2.ROTATE_90_CLOCKWISE)
        elif angle == 180:
            rotated = cv2.rotate(gray_img, cv2.ROTATE_180)
        else: # 270
            rotated = cv2.rotate(gray_img, cv2.ROTATE_90_COUNTERCLOCKWISE)

        score = pan_ocr_score(rotated)
        
        # SHORT-CIRCUIT: If we have extremely high confidence at 0 degrees, stop early.
        # A score > 1500 typically means multiple high-weight phrases (INCOME TAX, etc.) were found.
        if angle == 0 and score > 1500:
            return 0, gray_img

        # Tie-breaker: If scores are very close, prefer 0
        if angle == 0:
            score *= 1.25 # 25% bias towards 0 degrees for comparison
            
        if score > best_score:
            best_score = score
            best_angle = angle
            best_img = rotated

    return best_angle, best_img

def extract_aadhaar_num(df):
    """Extract Aadhaar number from Tesseract's DataFrame output."""
    aadhaar_candidates = df[
        (df['conf'].astype(int) >= 10) & (df['text'].str.match(r'^\d{4}$'))
    ].reset_index(drop=True)

    groups = []
    current_group = []

    for i in range(len(aadhaar_candidates)):
        row = aadhaar_candidates.iloc[i]
        if not current_group:
            current_group.append(row)
            continue

        prev = current_group[-1]
        same_line = abs(row['top'] - prev['top']) < 10
        close_x = row['left'] - (prev['left'] + prev['width']) < 40

        if same_line and close_x:
            current_group.append(row)
        else:
            if len(current_group) == 3:
                groups.append(current_group)
            current_group = [row]

    if len(current_group) == 3:
        groups.append(current_group)

    if groups:
        best_group = groups[0]
        aadhaar_num = "".join(item["text"] for item in best_group)
        # Bounding boxes for first 2 groups (8 digits) for masking
        mask_bboxes = [
            {"left": item["left"], "top": item["top"], "width": item["width"], "height": item["height"]}
            for item in best_group[:2]
        ]
        return aadhaar_num, mask_bboxes
    else:
        return None, None