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

# Explicitly set the path to the Tesseract executable.
# This is crucial in container environments where the system's PATH might not
# be fully inherited or visible to pytesseract in the way it expects by default.
# /usr/bin/tesseract is the standard path for Tesseract executable on Debian/Ubuntu based systems.
pytesseract.pytesseract.tesseract_cmd = r'/usr/bin/tesseract'


def preprocess_image(image: Image.Image) -> Image.Image:
    """
    Preprocesses the image for better OCR accuracy.
    Converts to grayscale, enhances contrast, and applies thresholding.
    """
    # Convert PIL Image to OpenCV format
    img_cv = np.array(image)
    if len(img_cv.shape) == 3: # Check if it's a color image
        img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2GRAY)

    # Apply adaptive thresholding
    # This often works better than simple thresholding for varying lighting conditions
    img_processed = cv2.adaptiveThreshold(
        img_cv, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
    )

    # Convert back to PIL Image
    return Image.fromarray(img_processed)


def detect_weight_and_units(text: str) -> list[str]:
    """
    Detects weight values and units from the extracted text.
    It looks for numbers followed by common weight units (case-insensitive).
    Examples: 70kg, 150 lbs, 50.5 g, 20.3kilograms
    """
    # Regex to find numbers (integers or decimals) followed by an optional space
    # and then one of the specified weight units.
    # The units are made case-insensitive with re.IGNORECASE.
    weight_patterns = re.findall(
        r'(\d+\.?\d*)\s*(kg|lbs|g|pounds|kilograms|grams)\b',
        text,
        re.IGNORECASE
    )

    detected_weights = []
    if weight_patterns:
        for value, unit in weight_patterns:
            detected_weights.append(f"{value} {unit.lower()}") # Standardize unit to lowercase
    return detected_weights


def process_image_for_weight(image: Image.Image) -> str:
    """
    Main function to process the uploaded image, perform OCR,
    and detect weight information.
    """
    if image is None:
        return "Please upload an image."

    try:
        # Preprocess the image for better OCR
        processed_img = preprocess_image(image)

        # Perform OCR on the preprocessed image
        extracted_text = pytesseract.image_to_string(processed_img)

        # Detect weight and units from the extracted text
        weights = detect_weight_and_units(extracted_text)

        if weights:
            # Join all detected weights into a single string
            return "Detected Weights:\n" + "\n".join(weights)
        else:
            return "No weight information detected. Extracted text:\n" + extracted_text

    except pytesseract.TesseractNotFoundError:
        # If this error still occurs, it means Tesseract isn't found at /usr/bin/tesseract
        # inside the Docker container, or there's an issue with the language data.
        return (
            "Tesseract OCR engine was not found at the specified path (/usr/bin/tesseract). "
            "This indicates a critical issue with the Docker environment setup. "
            "Please ensure the `Dockerfile` correctly installs 'tesseract-ocr' and 'tesseract-ocr-eng', "
            "and that the Tesseract executable is indeed located at `/usr/bin/tesseract` within the container."
        )
    except Exception as e:
        return f"An error occurred: {e}"


# Create the Gradio interface
iface = gr.Interface(
    fn=process_image_for_weight,
    inputs=gr.Image(type="pil", label="Upload Weight Log Image"),
    outputs=gr.Textbox(label="Detected Weight and Units"),
    title="Auto Weight Logger",
    description="Upload an image containing weight logs, and I will try to detect the weight values and their units.",
    allow_flagging="auto", # Allows users to flag inputs/outputs for debugging/improvement
    examples=[
        # You can add example images here if you have them.
        # For a live demo, these would be paths to actual images.
        # e.g., ["./examples/weight_log_1.png", "./examples/weight_log_2.jpg"]
    ]
)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()