import streamlit as st

import io
import pandas as pd
import numpy as np
import cv2
import pytesseract
from PIL import Image

from ultralytics import YOLO
from sahi import AutoDetectionModel
from sahi.predict import get_prediction, get_sliced_prediction
from sahi.utils.cv import visualize_object_predictions, read_image


def split_pdf_to_images(uploaded_file, read_status=None):
    """
    Takes a PDF file, converts each page to a PNG image, 
    and returns a list of byte objects.
    """
    try:
        # Convert PDF bytes to a list of PIL Image objects
        images = convert_from_bytes(uploaded_file.read())
        
        png_files = []
        for i, image in enumerate(images):
            
            if read_status is not None:
                try:
                    read_status.update(label=f"Downloading image {i+1}/{len(images)}")
                except Exception:
                    # DeltaGenerator may expose .text; fallback to that, then to st.write
                    try:
                        read_status.text(f"Downloading image {i+1}/{len(images)}")
                    except Exception:
                        st.write(f"Downloading Image: {i+1}/{len(images)}")
            else:
                st.write(f"Downloading Image: {i+1}/{len(images)}")
            # Save PIL image to a byte buffer as a PNG
            buf = io.BytesIO()
            image.save(buf, format="PNG")
            byte_im = buf.getvalue()
            png_files.append({"name": f"page_{i+1}.png", "content": byte_im})
            
        # Reset pointer for other functions
        uploaded_file.seek(0)
        return png_files
    except Exception as e:
        return f"Error processing PDF: {e}"


def analyze_data_types(uploaded_file):
    # (Keeping your second function for the other file drop)

    try:
        df = pd.read_csv(uploaded_file)
        uploaded_file.seek(0)
        return df.dtypes
    except Exception as e:
        return f"Error: {e}"


def run_computer_vision(image_list, analysis_status=None):
    """
    Runs CV model on images, draws bounding boxes, 
    and returns analysis with processed images.
    """
    results = []
    
    # Assuming 'model' is pre-loaded (e.g., model = YOLO('yolov8n.pt'))
    
    detection_model = AutoDetectionModel.from_pretrained(
                        model_type="ultralytics",
                        model_path="Model/Prod/10_epoch_model_fixed_text_size.pt",
                        confidence_threshold=0.3,
                        device="cpu",  # or 'cuda:0'
                    )
    
    results = []
    for i in range(len(image_list)):
        
        # Update the UI via the passed-in status placeholder when available.
        if analysis_status is not None:
            try:
                analysis_status.update(label=f"Processing image {i+1}/{len(image_list)}")
            except Exception:
                # DeltaGenerator may expose .text; fallback to that, then to st.write
                try:
                    analysis_status.text(f"Processing image {i+1}/{len(image_list)}")
                except Exception:
                    st.write(f"Processing Image: {i+1}/{len(image_list)}")
        else:
            st.write(f"Processing Image: {i+1}/{len(image_list)}")
        
        image_mem = Image.open(io.BytesIO(image_list[i]["content"])).convert("RGB")
        
        sahi_result = get_sliced_prediction(
        image_mem,
        detection_model,
        slice_height=256,
        slice_width=256,
        overlap_height_ratio=0.2,
        overlap_width_ratio=0.2,
        )
        
        object_prediction_list = sahi_result.object_prediction_list

        visualization_result = visualize_object_predictions(
            image=np.array(image_mem),
            object_prediction_list=object_prediction_list,
            hide_labels=False, # This removes the text labels entirely
            hide_conf=False,  # This removes the confidence scores
            rect_th=2         # Optional: Adjust box thickness
        )

        # 5. Convert the resulting numpy array back to a PIL Image
        annotated_canvas = visualization_result["image"]
        final_image = Image.fromarray(annotated_canvas)
        results.append(final_image.copy())
        
    return memory_images_to_pdf(results) # Returning the list of dicts containing image bytes and text


def memory_images_to_pdf(pil_image_list):
    if not pil_image_list:
        return None

    # Create an in-memory buffer for the PDF
    pdf_buffer = io.BytesIO()
    
    # Ensure all images are RGB
    rgb_images = [img.convert("RGB") for img in pil_image_list]
    
    # Save to the buffer
    rgb_images[0].save(
        pdf_buffer, 
        format="PDF", 
        save_all=True, 
        append_images=rgb_images[1:]
    )
    
    # Reset buffer position to the start so it can be read
    pdf_buffer.seek(0)
    return pdf_buffer

def get_codes_from_image():

    # *Optional: Set the path to the Tesseract executable if it's not in your system's PATH*
    # pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

    # Load the image using OpenCV
    image = cv2.imread('your_image.png')

    # Convert the image to grayscale (improves accuracy)
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Use Tesseract to extract text from the image
    text = pytesseract.image_to_string(gray_image)

    # Print the extracted text
    return text