"""
Crime Scene Analyzer Module
===========================
Provides a CrimeSceneAnalyzer class for forensic image analysis.
"""

import os
import json
import base64
import io
from datetime import datetime
from pathlib import Path
from dotenv import load_dotenv

import torch
from ultralytics import YOLO
from groq import Groq
from PIL import Image, ImageDraw, ImageFont

from reportlab.lib.pagesizes import A4
from reportlab.lib import colors
from reportlab.lib.units import mm
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.enums import TA_LEFT, TA_CENTER, TA_RIGHT, TA_JUSTIFY
from reportlab.platypus import (
    SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle,
    HRFlowable, Image as RLImage, KeepTogether
)
from reportlab.pdfgen import canvas
from reportlab.lib.utils import ImageReader

# ── palette ────────────────────────────────────────────────────────────────
C_BLACK      = colors.HexColor("#0D0D0D")
C_DARK_GRAY  = colors.HexColor("#1C1C1E")
C_MID_GRAY   = colors.HexColor("#3A3A3C")
C_LIGHT_GRAY = colors.HexColor("#AEAEB2")
C_WHITE      = colors.white
C_ACCENT     = colors.HexColor("#BF0000")
C_ACCENT2    = colors.HexColor("#E8A000")
C_GRID       = colors.HexColor("#2C2C2E")
C_ROW_ALT    = colors.HexColor("#1A1A1C")
C_HIGH       = colors.HexColor("#3A0000")
C_MED        = colors.HexColor("#3A2A00")
C_LOW        = colors.HexColor("#1A2A1A")

BOX_COLORS = [
    "#BF0000", "#E8A000", "#2060C0", "#20A060",
    "#8020C0", "#C06000", "#206080", "#C02060",
    "#60A000", "#804020",
]

# ── helpers ────────────────────────────────────────────────────────────────

def confidence_tier(conf: float) -> str:
    if conf >= 0.80:
        return "HIGH"
    if conf >= 0.50:
        return "MEDIUM"
    return "LOW"


# ── step 1 : detect objects via YOLOVIT ──────────────────────────────

def detect_objects_pil(img: Image.Image, model: YOLO) -> list[dict]:
    """Detect objects in a PIL Image using YOLO model."""
    print("[1/4] Detecting objects with YOLOVIT …")
    
    # Save image to temp file for YOLO
    temp_path = "temp_image.jpg"
    img.save(temp_path)
    
    # Run inference
    results = model(temp_path, conf=0.25, iou=0.45)
    
    # Get image dimensions
    W, H = img.size
    
    objects = []
    category_mapping = {
        'person': 'person',
        'car': 'vehicle', 'truck': 'vehicle', 'bus': 'vehicle', 'motorcycle': 'vehicle', 'bicycle': 'vehicle',
        'knife': 'weapon', 'gun': 'weapon', 'rifle': 'weapon', 'pistol': 'weapon',
        'cell phone': 'evidence', 'laptop': 'evidence', 'camera': 'evidence',
        'bottle': 'evidence', 'cup': 'evidence', 'wine glass': 'evidence',
        'chair': 'environmental', 'table': 'environmental', 'couch': 'environmental',
        'bed': 'environmental', 'tv': 'environmental',
        'book': 'document', 'paper': 'document', 'notebook': 'document',
    }
    
    for result in results:
        boxes = result.boxes
        for i, box in enumerate(boxes):
            # Get box coordinates in pixels
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            
            # Convert to normalized coordinates (0.0-1.0)
            x = x1 / W
            y = y1 / H
            w = (x2 - x1) / W
            h = (y2 - y1) / H
            
            # Get class label and confidence
            class_id = int(box.cls[0])
            confidence = float(box.conf[0])
            label = result.names[class_id]
            
            # Map to forensic category
            category = category_mapping.get(label.lower(), 'other')
            
            # Generate forensic notes based on label
            notes = f"Detected {label} with {confidence*100:.0f}% confidence"
            
            objects.append({
                "id": i + 1,
                "label": label.capitalize(),
                "category": category,
                "confidence": confidence,
                "bbox": {"x": x, "y": y, "w": w, "h": h},
                "notes": notes
            })
    
    # Clean up temp file
    if os.path.exists(temp_path):
        os.remove(temp_path)
    
    print(f"    → {len(objects)} objects detected")
    return objects


# ── step 2 : draw bounding boxes ───────────────────────────────────────────

def draw_bounding_boxes_pil(img: Image.Image, objects: list[dict]) -> Image.Image:
    """Draw bounding boxes on a PIL Image."""
    print("[2/4] Drawing bounding boxes …")
    img = img.convert("RGB")
    W, H = img.size
    draw = ImageDraw.Draw(img, "RGBA")

    # Font search: Linux → macOS → Windows, fall back to PIL default
    _font_candidates = [
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",       # Linux
        "/Library/Fonts/Arial Bold.ttf",                               # macOS
        "C:/Windows/Fonts/arialbd.ttf",                                # Windows
        "C:/Windows/Fonts/arial.ttf",                                  # Windows fallback
    ]
    def _load_font(size):
        for fp in _font_candidates:
            try:
                return ImageFont.truetype(fp, size)
            except Exception:
                continue
        return ImageFont.load_default()

    try:
        font_label = _load_font(max(12, W // 60))
        font_id    = _load_font(max(10, W // 70))
    except Exception:
        font_label = ImageFont.load_default()
        font_id    = font_label

    for i, obj in enumerate(objects):
        color_hex = BOX_COLORS[i % len(BOX_COLORS)]
        color_rgb = tuple(int(color_hex.lstrip("#")[j:j+2], 16) for j in (0, 2, 4))

        bbox = obj.get("bbox", {})
        x = int(bbox.get("x", 0) * W)
        y = int(bbox.get("y", 0) * H)
        w = int(bbox.get("w", 0.1) * W)
        h = int(bbox.get("h", 0.1) * H)
        x2, y2 = x + w, y + h

        # semi-transparent fill
        draw.rectangle([x, y, x2, y2], fill=(*color_rgb, 40), outline=(*color_rgb, 220), width=max(2, W // 300))

        # corner ticks
        tick = max(8, W // 80)
        for cx, cy, dx, dy in [(x, y, 1, 1), (x2, y, -1, 1), (x, y2, 1, -1), (x2, y2, -1, -1)]:
            draw.line([cx, cy, cx + dx * tick, cy], fill=color_rgb, width=3)
            draw.line([cx, cy, cx, cy + dy * tick], fill=color_rgb, width=3)

        # label badge
        label_text = f"[{obj['id']}] {obj['label']}"
        conf_text  = f"{obj['confidence']*100:.0f}%"
        bbox_l = draw.textbbox((0, 0), label_text, font=font_label)
        bbox_c = draw.textbbox((0, 0), conf_text, font=font_id)
        lw = max(bbox_l[2] - bbox_l[0], bbox_c[2] - bbox_c[0]) + 12
        lh = (bbox_l[3] - bbox_l[1]) + (bbox_c[3] - bbox_c[1]) + 10

        tag_y = max(0, y - lh - 4)
        draw.rectangle([x, tag_y, x + lw, tag_y + lh], fill=(*color_rgb, 220))
        draw.text((x + 6, tag_y + 2), label_text, font=font_label, fill=(255, 255, 255))
        draw.text((x + 6, tag_y + (bbox_l[3] - bbox_l[1]) + 4), conf_text, font=font_id, fill=(220, 220, 220))

    return img


# ── step 3 : scene summary via Groq ─────────────────────────────────────

SUMMARY_PROMPT = """You are a senior forensic analyst writing an official crime-scene report.

Based on the list of detected objects below, write a structured
forensic narrative summary. Cover:

1. SCENE OVERVIEW — general description of the environment
2. KEY FINDINGS — most significant pieces of evidence
3. PROBABLE SEQUENCE OF EVENTS — what may have happened
4. AREAS REQUIRING FURTHER INVESTIGATION
5. INITIAL RISK ASSESSMENT — any ongoing hazards

Detected objects:
{objects_json}

Write in formal, third-person forensic language. Be concise but thorough.
Total length: 250-400 words.
"""

def generate_summary_text(objects: list[dict], client: Groq) -> str:
    """Generate forensic summary using Groq API."""
    print("[3/4] Generating forensic summary with Groq …")
    
    obj_json = json.dumps([{"id": o["id"], "label": o["label"],
                             "category": o["category"], "confidence": o["confidence"],
                             "notes": o.get("notes", "")} for o in objects], indent=2)
    
    prompt = SUMMARY_PROMPT.format(objects_json=obj_json)
    
    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ],
        max_tokens=1024
    )
    return response.choices[0].message.content.strip()


# ── step 4 : build PDF report ──────────────────────────────────────────────

class ForensicReportCanvas:
    """Adds header/footer watermark to every page."""

    def __init__(self, case_id: str, analyst: str):
        self.case_id = case_id
        self.analyst = analyst

    def __call__(self, canv: canvas.Canvas, doc):
        W, H = A4
        canv.saveState()

        # ── top bar ──
        canv.setFillColor(C_BLACK)
        canv.rect(0, H - 18*mm, W, 18*mm, fill=1, stroke=0)
        canv.setFillColor(C_ACCENT)
        canv.rect(0, H - 19*mm, W, 1*mm, fill=1, stroke=0)

        canv.setFont("Helvetica-Bold", 7)
        canv.setFillColor(C_LIGHT_GRAY)
        canv.drawString(12*mm, H - 10*mm, "RESTRICTED — FORENSIC INTELLIGENCE UNIT")
        canv.drawRightString(W - 12*mm, H - 10*mm, f"CASE {self.case_id}")
        canv.setFont("Helvetica", 6)
        canv.drawRightString(W - 12*mm, H - 14*mm,
                             f"Analyst: {self.analyst}  |  {datetime.now().strftime('%Y-%m-%d %H:%M')}")

        # ── bottom bar ──
        canv.setFillColor(C_BLACK)
        canv.rect(0, 0, W, 12*mm, fill=1, stroke=0)
        canv.setFillColor(C_ACCENT)
        canv.rect(0, 12*mm, W, 0.5*mm, fill=1, stroke=0)

        canv.setFont("Helvetica", 6)
        canv.setFillColor(C_LIGHT_GRAY)
        canv.drawString(12*mm, 4*mm, "CONFIDENTIAL — NOT FOR PUBLIC RELEASE")
        canv.drawCentredString(W/2, 4*mm, f"Page {doc.page}")
        canv.drawRightString(W - 12*mm, 4*mm, "AI-ASSISTED ANALYSIS")

        canv.restoreState()


def pil_to_rl_image(pil_img: Image.Image, max_w_mm: float, max_h_mm: float) -> RLImage:
    buf = io.BytesIO()
    pil_img.save(buf, format="PNG")
    buf.seek(0)
    W_px, H_px = pil_img.size
    aspect = H_px / W_px
    w = min(max_w_mm * mm, (max_h_mm / aspect) * mm)
    h = w * aspect
    if h > max_h_mm * mm:
        h = max_h_mm * mm
        w = h / aspect
    return RLImage(buf, width=w, height=h)


def build_pdf_bytes(
    img: Image.Image,
    annotated_img: Image.Image,
    objects: list[dict],
    summary: str,
    case_id: str,
    analyst: str = "AI System v1.0"
) -> bytes:
    """Build PDF report and return as bytes."""
    print("[4/4] Building forensic PDF report …")
    
    buf = io.BytesIO()
    doc = SimpleDocTemplate(
        buf,
        pagesize=A4,
        topMargin=22*mm,
        bottomMargin=16*mm,
        leftMargin=15*mm,
        rightMargin=15*mm,
    )

    W, H = A4
    usable_w = W - 30*mm

    # ── styles ──
    ss = getSampleStyleSheet()

    def style(name, parent="Normal", **kw) -> ParagraphStyle:
        return ParagraphStyle(name, parent=ss[parent], **kw)

    S = {
        "title": style("title", fontSize=22, fontName="Helvetica-Bold",
                       textColor=C_WHITE, spaceAfter=2, leading=26),
        "subtitle": style("subtitle", fontSize=9, fontName="Helvetica",
                          textColor=C_ACCENT, spaceAfter=6, leading=12),
        "section": style("section", fontSize=11, fontName="Helvetica-Bold",
                         textColor=C_ACCENT, spaceBefore=10, spaceAfter=4),
        "meta_label": style("meta_label", fontSize=7, fontName="Helvetica-Bold",
                            textColor=C_LIGHT_GRAY),
        "meta_value": style("meta_value", fontSize=8, fontName="Helvetica",
                            textColor=C_WHITE),
        "body": style("body", fontSize=8.5, fontName="Helvetica",
                      textColor=colors.HexColor("#D0D0D0"), leading=13,
                      spaceAfter=6, alignment=TA_JUSTIFY),
        "caption": style("caption", fontSize=7, fontName="Helvetica",
                         textColor=C_LIGHT_GRAY, alignment=TA_CENTER, spaceAfter=4),
        "tbl_hdr": style("tbl_hdr", fontSize=7.5, fontName="Helvetica-Bold",
                         textColor=C_WHITE),
        "tbl_cell": style("tbl_cell", fontSize=7.5, fontName="Helvetica",
                          textColor=colors.HexColor("#CCCCCC"), leading=10),
        "tbl_notes": style("tbl_notes", fontSize=7, fontName="Helvetica-Oblique",
                           textColor=C_LIGHT_GRAY, leading=9),
        "badge_high": style("badge_high", fontSize=7, fontName="Helvetica-Bold",
                            textColor=colors.HexColor("#FF6060")),
        "badge_med":  style("badge_med",  fontSize=7, fontName="Helvetica-Bold",
                            textColor=colors.HexColor("#FFC060")),
        "badge_low":  style("badge_low",  fontSize=7, fontName="Helvetica-Bold",
                            textColor=colors.HexColor("#60FF60")),
    }

    story = []

    # ════════════════════════════════════════════════════════════════
    # COVER / TITLE BLOCK
    # ════════════════════════════════════════════════════════════════
    title_data = [[
        Paragraph("FORENSIC SCENE<br/>ANALYSIS REPORT", S["title"]),
        Paragraph(f"CASE NO.<br/><font size=14 color='#BF0000'><b>{case_id}</b></font>",
                  S["subtitle"])
    ]]
    title_tbl = Table(title_data, colWidths=[usable_w * 0.65, usable_w * 0.35])
    title_tbl.setStyle(TableStyle([
        ("BACKGROUND",    (0, 0), (-1, -1), C_BLACK),
        ("TOPPADDING",    (0, 0), (-1, -1), 10),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 10),
        ("LEFTPADDING",   (0, 0), (-1, -1), 10),
        ("RIGHTPADDING",  (0, 0), (-1, -1), 10),
        ("VALIGN",        (0, 0), (-1, -1), "MIDDLE"),
        ("ALIGN",         (1, 0), (1, 0),   "RIGHT"),
        ("LINEBELOW",     (0, 0), (-1, -1), 1, C_ACCENT),
    ]))
    story.append(title_tbl)
    story.append(Spacer(1, 4*mm))

    # ── meta block ──
    ts = datetime.now()
    meta_rows = [
        ["SOURCE IMAGE",  "uploaded_image.jpg",
         "ANALYSIS DATE", ts.strftime("%d %B %Y")],
        ["ANALYST",       analyst,
         "TIME",          ts.strftime("%H:%M:%S UTC")],
        ["TOTAL OBJECTS", str(len(objects)),
         "CLASSIFICATION", "RESTRICTED"],
    ]
    meta_col_w = [usable_w * 0.17, usable_w * 0.33, usable_w * 0.17, usable_w * 0.33]
    meta_tbl_data = []
    for row in meta_rows:
        meta_tbl_data.append([
            Paragraph(row[0], S["meta_label"]),
            Paragraph(row[1], S["meta_value"]),
            Paragraph(row[2], S["meta_label"]),
            Paragraph(row[3], S["meta_value"]),
        ])
    meta_tbl = Table(meta_tbl_data, colWidths=meta_col_w)
    meta_tbl.setStyle(TableStyle([
        ("BACKGROUND",    (0, 0), (-1, -1), C_DARK_GRAY),
        ("ROWBACKGROUNDS",(0, 0), (-1, -1), [C_DARK_GRAY, C_MID_GRAY]),
        ("TOPPADDING",    (0, 0), (-1, -1), 4),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 4),
        ("LEFTPADDING",   (0, 0), (-1, -1), 8),
        ("RIGHTPADDING",  (0, 0), (-1, -1), 8),
        ("LINEBELOW",     (0, -1), (-1, -1), 0.5, C_ACCENT),
    ]))
    story.append(meta_tbl)
    story.append(Spacer(1, 6*mm))

    # ════════════════════════════════════════════════════════════════
    # ANNOTATED IMAGE
    # ════════════════════════════════════════════════════════════════
    story.append(Paragraph("SECTION 1 — ANNOTATED SCENE IMAGE", S["section"]))
    story.append(HRFlowable(width="100%", thickness=0.5, color=C_ACCENT))
    story.append(Spacer(1, 3*mm))

    rl_img = pil_to_rl_image(annotated_img, max_w_mm=180, max_h_mm=120)
    img_tbl = Table([[rl_img]], colWidths=[usable_w])
    img_tbl.setStyle(TableStyle([
        ("BACKGROUND",    (0, 0), (-1, -1), C_BLACK),
        ("ALIGN",         (0, 0), (-1, -1), "CENTER"),
        ("TOPPADDING",    (0, 0), (-1, -1), 6),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 6),
        ("LINEBELOW",     (0, 0), (-1, -1), 1, C_ACCENT),
    ]))
    story.append(img_tbl)
    story.append(Paragraph(
        f"Figure 1.1 — Annotated crime scene image. "
        f"{len(objects)} objects identified and marked with coloured bounding boxes.",
        S["caption"]))
    story.append(Spacer(1, 5*mm))

    # ════════════════════════════════════════════════════════════════
    # EVIDENCE TABLE
    # ════════════════════════════════════════════════════════════════
    story.append(Paragraph("SECTION 2 — DETECTED OBJECTS & EVIDENCE INVENTORY", S["section"]))
    story.append(HRFlowable(width="100%", thickness=0.5, color=C_ACCENT))
    story.append(Spacer(1, 3*mm))

    col_w = [
        usable_w * 0.05,   # ID
        usable_w * 0.22,   # Label
        usable_w * 0.14,   # Category
        usable_w * 0.11,   # Confidence
        usable_w * 0.11,   # Tier
        usable_w * 0.37,   # Notes
    ]
    tbl_data = [[
        Paragraph("ID",         S["tbl_hdr"]),
        Paragraph("LABEL",      S["tbl_hdr"]),
        Paragraph("CATEGORY",   S["tbl_hdr"]),
        Paragraph("CONFIDENCE", S["tbl_hdr"]),
        Paragraph("TIER",       S["tbl_hdr"]),
        Paragraph("FORENSIC NOTES", S["tbl_hdr"]),
    ]]

    tier_style_map = {"HIGH": S["badge_high"], "MEDIUM": S["badge_med"], "LOW": S["badge_low"]}
    tier_bg_map    = {"HIGH": C_HIGH, "MEDIUM": C_MED, "LOW": C_LOW}

    tbl_row_bgs = []
    for i, obj in enumerate(objects):
        tier = confidence_tier(obj["confidence"])
        row_bg = C_ROW_ALT if i % 2 == 0 else C_DARK_GRAY
        tbl_row_bgs.append(row_bg)
        tbl_data.append([
            Paragraph(str(obj["id"]),                     S["tbl_cell"]),
            Paragraph(obj["label"],                       S["tbl_cell"]),
            Paragraph(obj.get("category", "other").upper(), S["tbl_cell"]),
            Paragraph(f"{obj['confidence']*100:.1f}%",   S["tbl_cell"]),
            Paragraph(tier,                               tier_style_map[tier]),
            Paragraph(obj.get("notes", "—"),             S["tbl_notes"]),
        ])

    ev_tbl = Table(tbl_data, colWidths=col_w, repeatRows=1)
    ts_cmds = [
        ("BACKGROUND",    (0, 0), (-1, 0),  C_ACCENT),
        ("TEXTCOLOR",     (0, 0), (-1, 0),  C_WHITE),
        ("TOPPADDING",    (0, 0), (-1, -1), 4),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 4),
        ("LEFTPADDING",   (0, 0), (-1, -1), 5),
        ("RIGHTPADDING",  (0, 0), (-1, -1), 5),
        ("ROWBACKGROUNDS",(0, 1), (-1, -1), [C_ROW_ALT, C_DARK_GRAY]),
        ("GRID",          (0, 0), (-1, -1), 0.3, C_GRID),
        ("VALIGN",        (0, 0), (-1, -1), "TOP"),
    ]
    ev_tbl.setStyle(TableStyle(ts_cmds))
    story.append(ev_tbl)
    story.append(Spacer(1, 5*mm))

    # ── category summary mini-table ──
    from collections import Counter
    cat_counts = Counter(o.get("category", "other") for o in objects)
    cat_rows = [[Paragraph("CATEGORY", S["tbl_hdr"]),
                 Paragraph("COUNT", S["tbl_hdr"]),
                 Paragraph("% OF TOTAL", S["tbl_hdr"])]]
    for cat, cnt in sorted(cat_counts.items(), key=lambda x: -x[1]):
        cat_rows.append([
            Paragraph(cat.upper(), S["tbl_cell"]),
            Paragraph(str(cnt), S["tbl_cell"]),
            Paragraph(f"{cnt/len(objects)*100:.1f}%", S["tbl_cell"]),
        ])
    cat_tbl = Table(cat_rows, colWidths=[usable_w*0.5, usable_w*0.25, usable_w*0.25])
    cat_tbl.setStyle(TableStyle([
        ("BACKGROUND",    (0, 0), (-1, 0),  C_MID_GRAY),
        ("BACKGROUND",    (0, 1), (-1, -1), C_DARK_GRAY),
        ("GRID",          (0, 0), (-1, -1), 0.3, C_GRID),
        ("TOPPADDING",    (0, 0), (-1, -1), 3),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 3),
        ("LEFTPADDING",   (0, 0), (-1, -1), 6),
        ("RIGHTPADDING",  (0, 0), (-1, -1), 6),
    ]))
    story.append(Paragraph("Table 2.2 — Category Distribution Summary", S["caption"]))
    story.append(cat_tbl)
    story.append(Spacer(1, 5*mm))

    # ════════════════════════════════════════════════════════════════
    # FORENSIC NARRATIVE SUMMARY
    # ════════════════════════════════════════════════════════════════
    story.append(Paragraph("SECTION 3 — FORENSIC NARRATIVE SUMMARY", S["section"]))
    story.append(HRFlowable(width="100%", thickness=0.5, color=C_ACCENT))
    story.append(Spacer(1, 3*mm))

    # parse the summary into sections
    for line in summary.split("\n"):
        line = line.strip()
        if not line:
            story.append(Spacer(1, 2*mm))
            continue
        # bold numbered headings
        if line and line[0].isdigit() and "." in line[:3]:
            story.append(Paragraph(
                f'<font color="#BF0000"><b>{line}</b></font>', S["body"]))
        else:
            story.append(Paragraph(line, S["body"]))

    story.append(Spacer(1, 5*mm))

    # ════════════════════════════════════════════════════════════════
    # DISCLAIMER
    # ════════════════════════════════════════════════════════════════
    story.append(HRFlowable(width="100%", thickness=0.5, color=C_MID_GRAY))
    story.append(Spacer(1, 2*mm))
    disclaimer = (
        "DISCLAIMER: This report was generated by an AI-assisted forensic analysis system. "
        "All findings are preliminary and must be reviewed and validated by a certified forensic "
        "investigator before being used in any legal or official capacity. Confidence scores are "
        "model estimates and do not constitute legal certainty."
    )
    story.append(Paragraph(disclaimer, style(
        "disclaimer", fontSize=6.5, textColor=C_LIGHT_GRAY,
        fontName="Helvetica-Oblique", alignment=TA_JUSTIFY)))

    # ── build ──
    page_fn = ForensicReportCanvas(case_id, analyst)
    doc.build(story, onFirstPage=page_fn, onLaterPages=page_fn)
    
    buf.seek(0)
    return buf.read()


# ── CrimeSceneAnalyzer Class ───────────────────────────────────────────

class CrimeSceneAnalyzer:
    """Main class for crime scene analysis."""
    
    def __init__(self):
        """Initialize the analyzer with models and API clients."""
        # Load environment variables
        load_dotenv()
        
        # Get Groq API key
        groq_api_key = os.environ.get("GROQ_API_KEY")
        if not groq_api_key:
            raise ValueError("GROQ_API_KEY not found in .env file or environment variables.")
        
        # Initialize Groq client
        self.groq_client = Groq(api_key=groq_api_key)
        
        # Load YOLO model
        print("Loading YOLOVIT model...")
        self.yolo_model = YOLO('yolov8x.pt')
        print("YOLOVIT model loaded successfully.")
    
    def analyze_scene(self, image: Image.Image) -> dict:
        """
        Analyze a crime scene image and return PDF bytes.
        
        Args:
            image: PIL Image object
            
        Returns:
            dict with keys:
                - pdf_bytes: bytes of the generated PDF
                - error: error message if analysis failed
        """
        try:
            # Generate case ID
            case_id = f"CS-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
            
            print(f"\n{'='*55}")
            print(f"  AI CRIME SCENE ANALYZER — Case {case_id}")
            print(f"{'='*55}\n")
            
            # Step 1: Detect objects
            objects = detect_objects_pil(image, self.yolo_model)
            
            # Step 2: Draw bounding boxes
            annotated_img = draw_bounding_boxes_pil(image, objects)
            
            # Step 3: Generate summary
            summary = generate_summary_text(objects, self.groq_client)
            
            # Step 4: Build PDF
            pdf_bytes = build_pdf_bytes(image, annotated_img, objects, summary, case_id)
            
            print(f"\n{'='*55}")
            print(f"  COMPLETE — PDF generated")
            print(f"{'='*55}\n")
            
            return {"pdf_bytes": pdf_bytes}
            
        except Exception as e:
            print(f"Error during analysis: {str(e)}")
            return {"error": str(e)}