import cv2
import os
import numpy as np

from core.config import settings
from core.constants import SEVERITY_SCORE
from core.logger import setup_logger, logger

from inspection.inference import DefectInspector
from inspection.geometry import analyze_defect
from inspection.severity import classify_defect
from inspection.tracker import DefectTracker
from inspection.lifecycle import DefectLifecycleManager
from inspection.output_formatter import (
    format_inspection_output,
    persist_inspection,
)
from inspection.class_map import CLASS_MAP
from inspection.service import annotate_image

from agent.langgraph_agent import run_agent


setup_logger()


def main():

    inspector = DefectInspector(settings.MODEL_PATH)
    tracker = DefectTracker()
    lifecycle = DefectLifecycleManager(settings.MAX_MISSING_FRAMES)

    image_files = sorted([
        f for f in os.listdir(settings.IMAGE_FOLDER)
        if f.lower().endswith((".jpg", ".png", ".jpeg"))
    ])

    index = 0
    repeat_count = 0
    frame_count = 0
    last_agent_result = None
    display_cycles = 0

    logger.info("Starting AI Inspection System...")

    while True:

        img_path = os.path.join(settings.IMAGE_FOLDER, image_files[index])
        frame = cv2.imread(img_path)

        if frame is None:
            continue

        # ---------------- SIMULATION ----------------
        if repeat_count < settings.IMAGE_REPEAT:
            display_frame = frame.copy()
            repeat_count += 1
        else:
            if repeat_count < settings.IMAGE_REPEAT + settings.BLANK_FRAMES:
                display_frame = np.zeros_like(frame)
                repeat_count += 1
            else:
                index = (index + 1) % len(image_files)
                repeat_count = 0
                continue

        frame_count += 1
        detection_frame = frame
        roi_offset = (0, 0)

        # ---------------- DETECTION ----------------
        detections = inspector.inspect_image(detection_frame)
        defects = tracker.update(detections)

        annotated = display_frame.copy()
        processed_defects = []

        # ---------------- PROCESS DEFECTS ----------------
        for d in defects:

            geometry = analyze_defect(d["contour"], detection_frame.shape)

            decision = classify_defect(
                geometry["area_pixels"],
                geometry["length_pixels"],
                geometry["area_ratio"]
            )

            x, y, w, h = d["bbox"]
            draw_x = x + roi_offset[0]
            draw_y = y + roi_offset[1]
            draw_contour = d["contour"].astype(np.int32) + np.array(
                [[[roi_offset[0], roi_offset[1]]]],
                dtype=np.int32,
            )

            color = (
                (0, 0, 255) if decision["decision"] == "FAIL"
                else (0, 165, 255) if decision["decision"] == "REVIEW"
                else (0, 255, 0)
            )

            # Bounding box
            cv2.rectangle(annotated, (draw_x, draw_y), (draw_x + w, draw_y + h), color, 2)

            # ---------------- SEGMENTATION MASK (FIX) ----------------
            # Filled overlay
            overlay = annotated.copy()
            cv2.drawContours(overlay, [draw_contour], -1, color, -1)
            cv2.addWeighted(overlay, 0.3, annotated, 0.7, 0, annotated)

            # Outline
            cv2.drawContours(annotated, [draw_contour], -1, color, 2)

            severity_score = SEVERITY_SCORE[decision["severity"]]

            processed_defects.append({
                "type": CLASS_MAP.get(d["class_id"], "unknown"),
                "severity": decision["severity"],
                "area_ratio": round(geometry["area_ratio"], 5),
                "length": round(geometry["length_pixels"], 2),
                "bbox": (draw_x, draw_y, w, h),
                "severity_score": severity_score
            })

        # ---------------- LIFECYCLE ----------------
        finalized = lifecycle.update(processed_defects, frame_count)

        if finalized:
            logger.info(f"Finalized {len(finalized)} defect(s)")

            output = format_inspection_output(finalized, source="simulation")

            try:
                last_agent_result = run_agent(output)
                output["decision"] = last_agent_result["decision"]
                output["recommendation"] = last_agent_result["recommendation"]
                output["summary_text"] = last_agent_result["summary"]
                output["agent_mode"] = last_agent_result.get("agent_mode", "heuristic")
                output["agent_provider"] = last_agent_result.get("agent_provider", "Rule-Based Safety Engine")
                output["agent_model"] = last_agent_result.get("agent_model", "fallback")
                logger.info("AI report generated")

                persist_inspection(output)

                # Show result for limited time
                display_cycles = 10

            except Exception as e:
                logger.error(f"Agent failed: {e}")

        # ---------------- UI STATUS (FIX) ----------------
        if processed_defects:
            text = "Processing..."
            color = (255, 255, 0)

        elif last_agent_result and display_cycles > 0:
            text = f"FINAL: {last_agent_result['decision']}"
            color = (
                (0, 0, 255) if text.endswith("FAIL")
                else (0, 165, 255) if text.endswith("REVIEW")
                else (0, 255, 0)
            )
            display_cycles -= 1

        else:
            text = "Idle"
            color = (255, 255, 255)

        cv2.putText(
            annotated,
            text,
            (30, 100),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.8,
            color,
            2
        )

        cv2.imshow("Inspection System", annotated)

        key = cv2.waitKey(500) & 0xFF
        if key == 27 or key == ord('q'):
            logger.info("Exit signal received")
            break

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()