import gradio as gr
import re
import pandas as pd
import os
import cv2
import numpy as np
from PIL import Image
import pytesseract

# Disable the DecompressionBomb limit to allow processing of extremely long screenshots
Image.MAX_IMAGE_PIXELS = None

def clean_and_format_number(phone_str):
    """
    Cleans the extracted string to ensure it's a valid phone number.
    Removes spaces, dashes, brackets, and names.
    """
    # Keep only digits and the plus sign
    cleaned = re.sub(r'[^\d+]', '', str(phone_str))
    
    # Extract just the digits to check length
    digits_only = cleaned.replace('+', '')
    
    # Most valid international numbers are between 10 and 15 digits long
    if 10 <= len(digits_only) <= 15:
        # Standardize format: ensure it starts with a '+' for Excel safety
        if not cleaned.startswith('+'):
            cleaned = '+' + cleaned
        return cleaned
    return None

def extract_numbers_from_images(image_files):
    """Main processing function triggered by Gradio."""
    if not image_files:
        return None, "⚠️ Please upload at least one image."

    all_unique_numbers = set()
    logs =[]

    for i, file_obj in enumerate(image_files):
        # Gradio 'file' type returns a string path or an object depending on the version.
        file_path = file_obj if isinstance(file_obj, str) else file_obj.name
        file_basename = os.path.basename(file_path)
        
        logs.append(f"⚡ Processing image {i+1}/{len(image_files)}: {file_basename}...")
        
        try:
            # 1. Read image incredibly fast using OpenCV
            # We use imdecode to safely handle file paths with special characters
            img_array = np.fromfile(file_path, np.uint8)
            img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
            
            if img is None:
                logs.append(f"❌ Error: Could not read image {file_basename}.")
                continue

            # 2. Convert to Grayscale (Speeds up OCR by 3x and improves contrast)
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            
            height, width = gray.shape
            
            # If the image is extremely wide, resize it down to 1080px width proportionally
            if width > 1080:
                ratio = 1080.0 / width
                new_height = int(height * ratio)
                gray = cv2.resize(gray, (1080, new_height), interpolation=cv2.INTER_AREA)
                height, width = gray.shape

            # 3. Slice image into safe chunks to prevent memory overload on massive screenshots
            MAX_HEIGHT = 4000
            chunk_count = 0
            
            for y in range(0, height, MAX_HEIGHT):
                chunk_count += 1
                # Crop the chunk using numpy slicing (instantaneous)
                chunk = gray[y:min(y + MAX_HEIGHT, height), :]
                
                # Run Tesseract OCR on the chunk
                text = pytesseract.image_to_string(chunk)
                
                # 4. Pure Logic Regex Extraction
                # Looks for a sequence of 9-25 chars containing digits, spaces, dashes, or brackets
                pattern = r'\+?[\d\s\-\(\)]{9,25}'
                potential_matches = re.findall(pattern, text)
                
                valid_count = 0
                for match in potential_matches:
                    clean_num = clean_and_format_number(match)
                    if clean_num:
                        all_unique_numbers.add(clean_num)
                        valid_count += 1
                
            logs.append(f"   -> Scanned {chunk_count} chunk(s). Current total unique numbers: {len(all_unique_numbers)}")
                    
        except Exception as e:
            logs.append(f"❌ Error processing {file_basename}: {str(e)}")

    # Check if we got any numbers
    if not all_unique_numbers:
        return None, "\n".join(logs) + "\n\n⚠️ No phone numbers were extracted from the uploaded images."

    # Convert set to sorted list
    sorted_numbers = sorted(list(all_unique_numbers))
    
    # Output only the phone numbers, as requested
    df = pd.DataFrame({"Phone Number": sorted_numbers})
    
    # Outputting to Excel to permanently prevent the Microsoft CSV Scientific Notation bug
    excel_filename = "extracted_whatsapp_numbers.xlsx"
    df.to_excel(excel_filename, index=False)
    
    final_status = "\n".join(logs) + f"\n\n🎉 SUCCESS! Instantly extracted {len(sorted_numbers)} total unique numbers."
    
    return excel_filename, final_status

# --- GRADIO INTERFACE ---

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.File(
                file_count="multiple", 
                type="filepath", 
                label="Upload WhatsApp Screenshots",
                file_types=["image"]
            )
            submit_btn = gr.Button("Extract Numbers Instantly", variant="primary")
            
        with gr.Column(scale=1):
            excel_output = gr.File(label="Download Extracted Numbers (Excel)")
            status_output = gr.Textbox(label="Processing Logs & Status", lines=12, interactive=False)
            
    submit_btn.click(
        fn=extract_numbers_from_images,
        inputs=[image_input],
        outputs=[excel_output, status_output]
    )

if __name__ == "__main__":
    demo.launch()