import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO
import os
from datetime import datetime
import matplotlib.pyplot as plt
from utils import create_percentage_circle, generate_pdf_report
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas

# Create directories with error handling
for directory in ["screenshots", "logs", "reports"]:
    try:
        os.makedirs(directory, exist_ok=True)
    except Exception as e:
        print(f"Warning: Could not create {directory} directory: {e}")

# Load YOLOv8 model
model = YOLO("yolov8m.pt")

def analyze_video(video_path, confidence=0.3):
    """
    Process video, detect objects, log results, capture screenshots, and generate PDF report.
    Args:
        video_path: Path to input video.
        confidence: Confidence threshold for detections.
    Returns:
        Tuple of (annotated video path, summary text, screenshot paths, percentage circle path, PDF path).
    """
    if not video_path:
        return None, "Error: No video uploaded.", [], None, None

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None, "Error: Could not open video.", [], None, None

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))

    output_path = "output_annotated.mp4"
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    log_file = f"logs/analysis_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
    vehicle_count = 0
    pedestrian_count = 0
    sign_count = 0
    alerts = []
    screenshots = []
    confidences = []

    with open(log_file, "w") as log:
        log.write(f"Analysis started at {datetime.now()}\n")
        log.write(f"Video: {video_path}\n")
        log.write(f"Confidence threshold: {confidence}\n\n")

        frame_count = 0
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            frame_count += 1

            results = model(frame, conf=confidence)
            frame_alerts = []
            detected_objects = []
            for result in results:
                boxes = result.boxes
                for box in boxes:
                    x1, y1, x2, y2 = map(int, box.xyxy[0])
                    cls = int(box.cls[0])
                    conf = float(box.conf[0])
                    label = f"{model.names[cls]} {conf:.2f}"
                    confidences.append(conf)
                    detected_objects.append((model.names[cls], conf))

                    if model.names[cls] == "person":
                        color = (0, 255, 0)
                    elif model.names[cls] in ["car", "truck", "bus"]:
                        color = (255, 0, 0)
                    else:
                        color = (0, 0, 255)
                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                    cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

                    if model.names[cls] in ["car", "truck", "bus"]:
                        vehicle_count += 1
                    elif model.names[cls] == "person":
                        pedestrian_count += 1
                    elif "sign" in model.names[cls].lower():
                        sign_count += 1

                    if model.names[cls] == "person":
                        for other_box in boxes:
                            other_cls = int(other_box.cls[0])
                            if model.names[other_cls] in ["car", "truck", "bus"]:
                                px1, py1, px2, py2 = map(int, box.xyxy[0])
                                vx1, vy1, vx2, vy2 = map(int, other_box.xyxy[0])
                                distance = np.sqrt((px1 - vx1)**2 + (py1 - vy1)**2)
                                if distance < 100:
                                    alert = f"Frame {frame_count}: Pedestrian near vehicle (distance: {distance:.1f}px, confidence: {conf:.2f})"
                                    frame_alerts.append(alert)
                                    screenshot_path = f"screenshots/frame_{frame_count}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
                                    cv2.imwrite(screenshot_path, frame)
                                    screenshots.append(screenshot_path)

                log.write(f"Frame {frame_count}: Detected {len(boxes)} objects\n")
                for obj, conf in detected_objects:
                    log.write(f"  - {obj} (confidence: {conf:.2f})\n")
                if frame_alerts:
                    log.write(f"  Alerts: {', '.join(frame_alerts)}\n")
                alerts.extend(frame_alerts)

            out.write(frame)

        cap.release()
        out.release()

        avg_confidence = np.mean(confidences) * 100 if confidences else 0
        circle_path = create_percentage_circle(avg_confidence, "Average Detection Confidence")

        summary = f"""
        Analysis Complete:
        - Total Frames Processed: {frame_count}
        - Average Vehicles per Frame: {vehicle_count / frame_count:.2f}
        - Average Pedestrians per Frame: {pedestrian_count / frame_count:.2f}
        - Traffic Signs Detected: {sign_count}
        - Safety Alerts: {len(alerts)}
        {chr(10).join(alerts) if alerts else "No critical incidents detected."}
        - Log File: {log_file}
        - Screenshots Saved: {len(screenshots)}
        """

        log.write(f"\nSummary:\n{summary}")

        # Generate PDF report
        pdf_path = f"reports/report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
        generate_pdf_report(summary, screenshots, pdf_path)

    return output_path, summary, screenshots, circle_path, pdf_path

css = """
.gradio-container {background-color: #f0f4f8; font-family: Arial;}
.footer {display: none !important;}
button {background-color: #27ae60; color: white; border-radius: 5px; padding: 10px 20px;}
button:hover {background-color: #219653;}
.gradio-gallery, .gradio-video, .gradio-textbox {border-radius: 5px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);}
"""

with gr.Blocks(css=css) as iface:
    gr.Markdown("# Road Safety AI Video Analysis")
    gr.Markdown("Upload a traffic video for analysis and download a detailed PDF report.")
    
    with gr.Row():
        video_input = gr.Video(label="Upload Video")
        confidence_input = gr.Slider(0.1, 1.0, value=0.3, step=0.1, label="Confidence Threshold")
    
    analyze_button = gr.Button("Analyze Video")
    
    with gr.Row():
        with gr.Column():
            video_output = gr.Video(label="Annotated Video")
            summary_output = gr.Textbox(label="Analysis Summary")
        with gr.Column():
            screenshot_output = gr.Gallery(label="Incident Screenshots")
            circle_output = gr.Image(label="Confidence Metric")
            pdf_output = gr.File(label="Download Report")
    
    analyze_button.click(
        fn=analyze_video,
        inputs=[video_input, confidence_input],
        outputs=[video_output, summary_output, screenshot_output, circle_output, pdf_output]
    )

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