# processor.py

import cv2
import numpy as np
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.mime.base import MIMEBase
from email import encoders
import os
from ultralytics import YOLO
from transformers import AutoModel, AutoProcessor
from PIL import Image, ImageDraw, ImageFont
import re
import torch

# Email credentials (replace with your App-Specific Password)
FROM_EMAIL = "Fares5675@gmail.com"
EMAIL_PASSWORD = "cawxqifzqiwjufde"  # Use App-Specific Password here
TO_EMAIL = "kduwaysan@gmail.com"
SMTP_SERVER = 'smtp.gmail.com'
SMTP_PORT = 465

# Arabic dictionary for converting license plate text
arabic_dict = {
    "0": "٠", "1": "١", "2": "٢", "3": "٣", "4": "٤", "5": "٥",
    "6": "٦", "7": "٧", "8": "٨", "9": "٩", "A": "ا", "B": "ب",
    "J": "ح", "D": "د", "R": "ر", "S": "س", "X": "ص", "T": "ط",
    "E": "ع", "G": "ق", "K": "ك", "L": "ل", "Z": "م", "N": "ن",
    "H": "ه", "U": "و", "V": "ي", " ": " "
}
class_colors = {
    0: (0, 255, 0),    # Green (Helmet)
    1: (255, 0, 0),    # Blue (License Plate)
    2: (0, 0, 255),    # Red (MotorbikeDelivery)
    3: (255, 255, 0),  # Cyan (MotorbikeSport)
    4: (255, 0, 255),  # Magenta (No Helmet)
    5: (0, 255, 255),  # Yellow (Person)
}

# Load the OCR model
processor = AutoProcessor.from_pretrained("stepfun-ai/GOT-OCR2_0", trust_remote_code=True)
model_ocr = AutoModel.from_pretrained("stepfun-ai/GOT-OCR2_0", trust_remote_code=True).to('cuda')

# Load YOLO model
model = YOLO('yolov8_Medium.pt')  # Update the path as needed

# Define lane area coordinates (example coordinates)
red_lane = np.array([[2,1583],[1,1131],[1828,1141],[1912,1580]], np.int32)

# Dictionary to track violations per license plate
violations_dict = {}

def filter_license_plate_text(license_plate_text):
    license_plate_text = re.sub(r'[^A-Z0-9]+', "", license_plate_text)
    match = re.search(r'(\d{4})\s*([A-Z]{2})', license_plate_text)
    return f"{match.group(1)} {match.group(2)}" if match else None

def convert_to_arabic(license_plate_text):
    return "".join(arabic_dict.get(char, char) for char in license_plate_text)

def send_email(license_text, violation_image_path, violation_type): 
    # Define the subject and body based on violation type
    if violation_type == 'No Helmet, In Red Lane':
        subject = 'تنبيه مخالفة: عدم ارتداء خوذة ودخول المسار الأيسر'
        body = f"لعدم ارتداء الخوذة ولدخولها المسار الأيسر ({license_text}) تم تغريم دراجة نارية التي تحمل لوحة"
        
    elif violation_type == 'In Red Lane':
        subject = 'تنبيه مخالفة: دخول المسار الأيسر'
        body = f"لدخولها المسار الأيسر ({license_text}) تم تغريم دراجة نارية التي تحمل لوحة"
    else:
        violation_type == 'No Helmet'
        subject = 'تنبيه مخالفة: عدم ارتداء خوذة'
        body = f"لعدم ارتداء الخوذة ({license_text}) تم تغريم دراجة نارية التي تحمل لوحة"
        
    # Create the email message
    msg = MIMEMultipart()
    msg['From'] = FROM_EMAIL
    msg['To'] = TO_EMAIL
    msg['Subject'] = subject 
    msg.attach(MIMEText(body, 'plain'))

    # Attach the violation image
    if os.path.exists(violation_image_path):
        with open(violation_image_path, 'rb') as attachment_file:
            part = MIMEBase('application', 'octet-stream')
            part.set_payload(attachment_file.read())
            encoders.encode_base64(part)
            part.add_header('Content-Disposition', f'attachment; filename={os.path.basename(violation_image_path)}')
            msg.attach(part)

    # Send the email using SMTP
    try:
        with smtplib.SMTP_SSL(SMTP_SERVER, SMTP_PORT) as server:
            server.login(FROM_EMAIL, EMAIL_PASSWORD)
            server.sendmail(FROM_EMAIL, TO_EMAIL, msg.as_string())
            print("Email with attachment sent successfully!")
    except Exception as e:
        print(f"Failed to send email: {e}")

def draw_text_pil(img, text, position, font_path, font_size, color):
    img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    
    draw = ImageDraw.Draw(img_pil)
    
    try:
        font = ImageFont.truetype(font_path, size=font_size)
    except IOError:
        print(f"Font file not found at {font_path}. Using default font.")
        font = ImageFont.load_default()
        
    draw.text(position, text, font=font, fill=color)
    
    img_np = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
    return img_np

def process_video(video_path, font_path, violation_image_path='violation.jpg'):
    cap = cv2.VideoCapture(video_path)
    
    if not cap.isOpened():
        print("Error opening video file")
        return None
    
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    output_video_path = 'output_violation.mp4'
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
    
    margin_y = 50
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break  

        cv2.polylines(frame, [red_lane], isClosed=True, color=(0, 0, 255), thickness=3)  

        results = model.track(frame)
        
        for box in results[0].boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())  
            label = model.names[int(box.cls)]  
            color = class_colors.get(int(box.cls), (255, 255, 255))
            confidence = box.conf[0].item()

            helmet_violation = False
            lane_violation = False
            violation_type = []

            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3)  
            cv2.putText(frame, f'{label}: {confidence:.2f}', (x1, y1 - 10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

            if label == 'MotorbikeDelivery' and confidence >= 0.4:
                motorbike_crop = frame[max(0, y1 - margin_y):y2, x1:x2]
                delivery_center = ((x1 + x2) // 2, y2)
                in_red_lane = cv2.pointPolygonTest(red_lane, delivery_center, False)
                if in_red_lane >= 0:
                    lane_violation = True
                    violation_type.append("In Red Lane")

                sub_results = model(motorbike_crop)

                for result in sub_results[0].boxes:
                    sub_x1, sub_y1, sub_x2, sub_y2 = map(int, result.xyxy[0].cpu().numpy())  
                    sub_label = model.names[int(result.cls)]
                    sub_color = (255, 0, 0)  

                    cv2.rectangle(motorbike_crop, (sub_x1, sub_y1), (sub_x2, sub_y2), sub_color, 2)
                    cv2.putText(motorbike_crop, sub_label, (sub_x1, sub_y1 - 10), 
                                cv2.FONT_HERSHEY_SIMPLEX, 0.6, sub_color, 2)

                    if sub_label == 'No_Helmet':
                        helmet_violation = True
                        violation_type.append("No Helmet")
                        continue
                    if sub_label == 'License_plate':
                        license_crop = motorbike_crop[sub_y1:sub_y2, sub_x1:sub_x2]

                        if helmet_violation or lane_violation:
                            cv2.imwrite(violation_image_path, frame)
                            license_plate_pil = Image.fromarray(cv2.cvtColor(license_crop, cv2.COLOR_BGR2RGB))
                            temp_image_path = 'license_plate.png'
                            license_plate_pil.save(temp_image_path)
                            # Placeholder for OCR
                            # license_plate_text = model_ocr.chat(processor, temp_image_path, ocr_type='ocr')
                            # For demonstration, we'll mock the OCR result
                            license_plate_text = "1234AB"
                            filtered_text = filter_license_plate_text(license_plate_text)

                            if filtered_text:
                                if filtered_text not in violations_dict:
                                    violations_dict[filtered_text] = violation_type  
                                    send_email(filtered_text, violation_image_path, ', '.join(violation_type)) 
                                else:
                                    current_violations = set(violations_dict[filtered_text]) 
                                    new_violations = set(violation_type) 
                                    updated_violations = list(current_violations | new_violations) 

                                    if updated_violations != violations_dict[filtered_text]:
                                        violations_dict[filtered_text] = updated_violations
                                        send_email(filtered_text, violation_image_path, ', '.join(updated_violations)) 

                                arabic_text = convert_to_arabic(filtered_text)
                                frame = draw_text_pil(frame, filtered_text, (x1, y2 + 30), font_path, font_size=30, color=(255, 255, 255))
                                frame = draw_text_pil(frame, arabic_text, (x1, y2 + 60), font_path, font_size=30, color=(0, 255, 0))

        out.write(frame)
    
    cap.release()
    out.release()
    return output_video_path

def process_image(image_path, font_path, violation_image_path='violation.jpg'):
    frame = cv2.imread(image_path)
    if frame is None:
        print("Error loading image")
        return None

    cv2.polylines(frame, [red_lane], isClosed=True, color=(0, 0, 255), thickness=3)  

    results = model.track(frame)
    
    for box in results[0].boxes:
        x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())  
        label = model.names[int(box.cls)]  
        color = class_colors.get(int(box.cls), (255, 255, 255))
        confidence = box.conf[0].item()

        helmet_violation = False
        lane_violation = False
        violation_type = []

        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3)  
        cv2.putText(frame, f'{label}: {confidence:.2f}', (x1, y1 - 10), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

        if label == 'MotorbikeDelivery' and confidence >= 0.4:
            motorbike_crop = frame[max(0, y1 - 50):y2, x1:x2]
            delivery_center = ((x1 + x2) // 2, y2)
            in_red_lane = cv2.pointPolygonTest(red_lane, delivery_center, False)
            if in_red_lane >= 0:
                lane_violation = True
                violation_type.append("In Red Lane")

            sub_results = model(motorbike_crop)

            for result in sub_results[0].boxes:
                sub_x1, sub_y1, sub_x2, sub_y2 = map(int, result.xyxy[0].cpu().numpy())  
                sub_label = model.names[int(result.cls)]
                sub_color = (255, 0, 0)  

                cv2.rectangle(motorbike_crop, (sub_x1, sub_y1), (sub_x2, sub_y2), sub_color, 2)
                cv2.putText(motorbike_crop, sub_label, (sub_x1, sub_y1 - 10), 
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, sub_color, 2)

                if sub_label == 'No_Helmet':
                    helmet_violation = True
                    violation_type.append("No Helmet")
                    continue
                if sub_label == 'License_plate':
                    license_crop = motorbike_crop[sub_y1:sub_y2, sub_x1:sub_x2]

                    if helmet_violation or lane_violation:
                        cv2.imwrite(violation_image_path, frame)
                        license_plate_pil = Image.fromarray(cv2.cvtColor(license_crop, cv2.COLOR_BGR2RGB))
                        temp_image_path = 'license_plate.png'
                        license_plate_pil.save(temp_image_path)
                        # Placeholder for OCR
                        # license_plate_text = model_ocr.chat(processor, temp_image_path, ocr_type='ocr')
                        # For demonstration, we'll mock the OCR result
                        license_plate_text = "1234AB"
                        filtered_text = filter_license_plate_text(license_plate_text)

                        if filtered_text:
                            if filtered_text not in violations_dict:
                                violations_dict[filtered_text] = violation_type  
                                send_email(filtered_text, violation_image_path, ', '.join(violation_type)) 
                            else:
                                current_violations = set(violations_dict[filtered_text]) 
                                new_violations = set(violation_type) 
                                updated_violations = list(current_violations | new_violations) 

                                if updated_violations != violations_dict[filtered_text]:
                                    violations_dict[filtered_text] = updated_violations
                                    send_email(filtered_text, violation_image_path, ', '.join(updated_violations)) 

                            arabic_text = convert_to_arabic(filtered_text)
                            frame = draw_text_pil(frame, filtered_text, (x1, y2 + 30), font_path, font_size=30, color=(255, 255, 255))
                            frame = draw_text_pil(frame, arabic_text, (x1, y2 + 60), font_path, font_size=30, color=(0, 255, 0))

    return frame