agents/answer_parser.py

"""
Parse raw model text output into structured data per agent type.

Each agent generates text — this module extracts structured fields
(detections, counts, alerts, text regions, reasoning chains) from that text.
"""

import re
from typing import Any


class AnswerParser:
    """Parse model text into structured fields for each agent type."""

    @staticmethod
    def parse(text: str, agent_type: str) -> dict[str, Any]:
        """Parse text based on agent type. Returns structured dict."""
        parsers = {
            "vqa": AnswerParser._parse_vqa,
            "detect": AnswerParser._parse_detect,
            "alert": AnswerParser._parse_alert,
            "caption": AnswerParser._parse_caption,
            "track": AnswerParser._parse_track,
            "count": AnswerParser._parse_count,
            "ocr": AnswerParser._parse_ocr,
            "reason": AnswerParser._parse_reason,
        }
        parser = parsers.get(agent_type, AnswerParser._parse_vqa)
        return parser(text)

    @staticmethod
    def _parse_vqa(text: str) -> dict:
        return {"output_type": "text"}

    @staticmethod
    def _parse_detect(text: str) -> dict:
        detections = []
        # Try "label at [x1, y1, x2, y2]" format
        bbox_pattern = re.findall(r'(\w[\w\s]*?)\s+at\s+\[([0-9.,\s]+)\]', text)
        if bbox_pattern:
            for label, coords in bbox_pattern:
                try:
                    bbox = [float(c.strip()) for c in coords.split(",")]
                    detections.append({"label": label.strip().lower(), "bbox": bbox, "confidence": 0.8})
                except ValueError:
                    pass
        else:
            # Simple comma-separated object list: "person, car, truck"
            objects = [o.strip() for o in text.replace(" and ", ", ").split(",") if o.strip()]
            for obj in objects:
                obj_clean = re.sub(r'[^a-zA-Z\s]', '', obj).strip().lower()
                if obj_clean and len(obj_clean) > 1:
                    detections.append({"label": obj_clean, "bbox": [], "confidence": 0.5})
        return {"output_type": "text+image", "detections": detections}

    @staticmethod
    def _parse_alert(text: str) -> dict:
        alert = {"severity": "LOW", "category": "unknown", "description": text}
        upper = text.upper()
        if "CRITICAL" in upper:
            alert["severity"] = "CRITICAL"
        elif "HIGH" in upper:
            alert["severity"] = "HIGH"
        elif "MEDIUM" in upper:
            alert["severity"] = "MEDIUM"
        for cat in ["zone_violation", "loitering", "fight", "fall", "fire", "theft",
                     "trespassing", "tailgating", "abandoned", "crowd"]:
            if cat.replace("_", " ") in text.lower() or cat in text.lower():
                alert["category"] = cat
                break
        return {"output_type": "text+image", "alert": alert}

    @staticmethod
    def _parse_caption(text: str) -> dict:
        attrs = {}
        for time in ["daytime", "nighttime", "dawn", "dusk", "morning", "evening"]:
            if time in text.lower():
                attrs["time_of_day"] = time
        for weather in ["overcast", "rainy", "sunny", "foggy", "cloudy", "clear"]:
            if weather in text.lower():
                attrs["weather"] = weather
        for loc in ["intersection", "parking", "lobby", "hallway", "street", "warehouse"]:
            if loc in text.lower():
                attrs["location"] = loc
        return {"output_type": "text+image", "scene_attributes": attrs}

    @staticmethod
    def _parse_track(text: str) -> dict:
        tracks = []
        direction = "unknown"
        for d in ["left to right", "right to left", "stationary", "entering", "exiting",
                   "approaching", "departing", "walking", "running"]:
            if d in text.lower():
                direction = d
                break
        tracks.append({"object_id": 1, "label": "object", "speed": direction, "trajectory": []})
        return {"output_type": "text+image+video", "tracks": tracks}

    @staticmethod
    def _parse_count(text: str) -> dict:
        counts = {}
        # Parse "3 people, 2 cars" patterns
        for match in re.finditer(r'(\d+)\s+(\w+)', text):
            num, obj = int(match.group(1)), match.group(2).lower().rstrip('s')
            counts[obj] = num
        # If just a bare number
        if not counts and text.strip().isdigit():
            counts["total"] = int(text.strip())
        if counts and "total" not in counts:
            counts["total"] = sum(counts.values())
        return {"output_type": "text+image", "counts": counts}

    @staticmethod
    def _parse_ocr(text: str) -> dict:
        text_regions = []
        for part in re.split(r'[,;]', text):
            part = part.strip()
            if part and len(part) > 1:
                rtype = "license_plate" if re.search(r'[A-Z]{2}\d{2}', part) else "text"
                text_regions.append({"text": part, "bbox": [], "type": rtype})
        return {"output_type": "text+image", "text_regions": text_regions}

    @staticmethod
    def _parse_reason(text: str) -> dict:
        analysis = {"risk_level": "LOW", "reasoning_chain": [], "recommendation": ""}
        upper = text.upper()
        if "CRITICAL" in upper or "immediate" in text.lower():
            analysis["risk_level"] = "HIGH"
        elif "HIGH" in upper or "concern" in text.lower() or "suspicious" in text.lower():
            analysis["risk_level"] = "MEDIUM"
        sentences = [s.strip() for s in re.split(r'[.!]', text) if s.strip()]
        analysis["reasoning_chain"] = sentences[:5]
        if sentences:
            analysis["recommendation"] = sentences[-1]
        return {"output_type": "text+image", "analysis": analysis}