from flask import Flask, request, render_template, send_file
from inference_sdk import InferenceHTTPClient
from PIL import Image, ImageDraw, ImageFont, ImageEnhance
import os
from collections import defaultdict

app = Flask(__name__)

# Ensure the "static/processed" directory exists for output images
PROCESSED_DIR = "static/processed"
if not os.path.exists(PROCESSED_DIR):
    os.makedirs(PROCESSED_DIR, exist_ok=True)

# Securely get API key from Hugging Face Secrets
API_KEY = os.getenv("ROBOFLOW_API_KEY")  

# Initialize the Roboflow client
CLIENT = InferenceHTTPClient(
    api_url="https://detect.roboflow.com",
    api_key=API_KEY  # Secure way to use API key
)

# Model settings
MODEL_ID = "hvacsym/5"
CONFIDENCE_THRESHOLD = 0.3  # Confidence threshold for filtering predictions
GRID_SIZE = (3, 3)  # 3x3 segmentation

# Colors for bounding boxes
RED = (255, 0, 0)
GREEN = (0, 255, 0)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)

# Load font for labeling
try:
    font = ImageFont.truetype("arial.ttf", 14)
except:
    font = ImageFont.load_default()

def enhance_image(image):
    """Enhance image by adjusting brightness and contrast."""
    if image.mode != 'L':
        image = image.convert('L')
    brightness = ImageEnhance.Brightness(image)
    image = brightness.enhance(1.3)
    contrast = ImageEnhance.Contrast(image)
    image = contrast.enhance(1.2)
    return image.convert('RGB')  # Convert back to RGB for colored boxes

def process_image(image_path):
    """Processes an image by running inference and drawing bounding boxes."""
    # Load and enhance the original image
    original_image = Image.open(image_path)
    original_image = enhance_image(original_image)
    width, height = original_image.size
    seg_w, seg_h = width // GRID_SIZE[1], height // GRID_SIZE[0]
    
    # Create a copy of the full image to draw bounding boxes
    final_image = original_image.copy()
    draw_final = ImageDraw.Draw(final_image)
    total_counts = defaultdict(int)
    
    # Process each segment
    for row in range(GRID_SIZE[0]):
        for col in range(GRID_SIZE[1]):
            x1, y1 = col * seg_w, row * seg_h
            x2, y2 = (col + 1) * seg_w, (row + 1) * seg_h
            segment = original_image.crop((x1, y1, x2, y2))
            segment_path = f"static/processed/segment_{row}_{col}.png"
            segment.save(segment_path)
            
            # Run inference on the segment
            result = CLIENT.infer(segment_path, model_id=MODEL_ID)
            
            # Filter predictions based on confidence
            filtered_predictions = [
                pred for pred in result["predictions"] if pred["confidence"] * 100 >= CONFIDENCE_THRESHOLD
            ]
            
            # Draw bounding boxes and count labels
            for obj in filtered_predictions:
                class_name = obj["class"]
                total_counts[class_name] += 1
                x_min, y_min = x1 + obj["x"] - obj["width"] // 2, y1 + obj["y"] - obj["height"] // 2
                x_max, y_max = x1 + obj["x"] + obj["width"] // 2, y1 + obj["y"] + obj["height"] // 2
                
                # Draw bounding box
                draw_final.rectangle([x_min, y_min, x_max, y_max], outline=GREEN, width=2)
                
                # Draw extended label above the bounding box
                text_size = draw_final.textbbox((0, 0), class_name, font=font)
                text_width = text_size[2] - text_size[0]
                text_height = text_size[3] - text_size[1]
                text_x = x_min
                text_y = y_min - text_height - 5 if y_min - text_height - 5 > 0 else y_max + 5
                
                draw_final.rectangle([text_x, text_y, text_x + text_width + 6, text_y + text_height + 2], fill=BLACK)
                draw_final.text((text_x + 3, text_y), class_name, fill=WHITE, font=font)
    
    # Save the final processed image
    final_image_path = "static/processed/processed_image.png"
    final_image.save(final_image_path)
    return final_image_path, total_counts

@app.route("/", methods=["GET", "POST"])
def index():
    if request.method == "POST":
        image_file = request.files["image"]
        if image_file:
            # Ensure static/processed directory exists
            os.makedirs(PROCESSED_DIR, exist_ok=True)

            image_path = os.path.join(PROCESSED_DIR, "uploaded_image.png")
            image_file.save(image_path)

            # Process image
            final_image_path, total_counts = process_image(image_path)

            return render_template("index.html", final_image=final_image_path, counts=total_counts)

    return render_template("index.html", final_image=None, counts=None)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)