import cv2
import numpy as np
import os
import sys
import json
from collections import Counter
from inference_sdk import InferenceHTTPClient

# ─────────────────────────────────────────────
#  CONFIGURATION
# ─────────────────────────────────────────────
ROBOFLOW_API_KEY  = "bVPbU8TisRCASiURr0lb"
MODEL_ID          = "printed-circuit-board/3"
CONFIDENCE_THRESH = 0.3
OVERLAP_THRESH    = 0.3

LABEL_COLORS = {
    "resistor":       (0,   255, 0),
    "capacitor":      (255, 0,   0),
    "inductor":       (0,   0,   255),
    "diode":          (255, 255, 0),
    "led":            (0,   255, 255),
    "ic":             (255, 0,   255),
    "transistor":     (128, 255, 0),
    "connector":      (0,   128, 255),
    "jumper":         (255, 128, 0),
    "emi_filter":     (128, 0,   255),
    "button":         (0,   255, 128),
    "clock":          (255, 0,   128),
    "transformer":    (128, 128, 0),
    "potentiometer":  (0,   128, 128),
    "heatsink":       (128, 0,   128),
    "fuse":           (200, 200, 0),
    "ferrite_bead":   (0,   200, 200),
    "buzzer":         (200, 0,   200),
    "display":        (100, 200, 255),
    "battery":        (255, 200, 100),
}
DEFAULT_COLOR = (255, 255, 255)


def get_client():
    return InferenceHTTPClient(
        api_url="https://serverless.roboflow.com",
        api_key=ROBOFLOW_API_KEY
    )


def run_detection(image_path):
    print(f"[->] Sending image to Roboflow API...")
    client = get_client()
    result = client.infer(image_path, model_id=MODEL_ID)

    detections = []
    predictions = result.get("predictions", [])
    img_w = result.get("image", {}).get("width", 1)
    img_h = result.get("image", {}).get("height", 1)

    for pred in predictions:
        confidence = pred.get("confidence", 0)
        if confidence < CONFIDENCE_THRESH:
            continue

        label = pred.get("class", "unknown").lower()
        cx = pred.get("x", 0)
        cy = pred.get("y", 0)
        w  = pred.get("width",  0)
        h  = pred.get("height", 0)

        x1 = max(0, int(cx - w / 2))
        y1 = max(0, int(cy - h / 2))
        x2 = min(img_w, int(cx + w / 2))
        y2 = min(img_h, int(cy + h / 2))

        detections.append({
            'label':      label,
            'confidence': round(confidence, 3),
            'bbox':       (x1, y1, x2, y2)
        })

    print(f"[OK] Roboflow returned {len(predictions)} predictions, "
          f"{len(detections)} above {CONFIDENCE_THRESH:.0%} confidence")
    return detections


def detect_traces(img):
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    lower_copper = np.array([10,  50,  50])
    upper_copper = np.array([30, 255, 255])
    copper_mask  = cv2.inRange(hsv, lower_copper, upper_copper)

    lower_silver = np.array([0, 0, 180])
    upper_silver = np.array([180, 30, 255])
    silver_mask  = cv2.inRange(hsv, lower_silver, upper_silver)

    trace_mask = cv2.bitwise_or(copper_mask, silver_mask)

    kernel = np.ones((2, 2), np.uint8)
    trace_mask = cv2.morphologyEx(trace_mask, cv2.MORPH_OPEN,  kernel, iterations=1)
    trace_mask = cv2.morphologyEx(trace_mask, cv2.MORPH_CLOSE, kernel, iterations=1)

    return trace_mask


def draw_detections(img, detections, trace_mask=None):
    output = img.copy()

    if trace_mask is not None:
        trace_overlay = np.zeros_like(output)
        trace_overlay[trace_mask > 0] = (255, 100, 0)
        output = cv2.addWeighted(output, 1.0, trace_overlay, 0.3, 0)

    for det in detections:
        x1, y1, x2, y2 = det['bbox']
        label           = det['label']
        conf            = det['confidence']
        color           = LABEL_COLORS.get(label, DEFAULT_COLOR)

        cv2.rectangle(output, (x1, y1), (x2, y2), color, 2)

        text = f"{label} {conf:.0%}"
        font       = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.45
        thickness  = 1
        (tw, th), _ = cv2.getTextSize(text, font, font_scale, thickness)
        cv2.rectangle(output,
                      (x1, max(0, y1 - th - 8)),
                      (x1 + tw + 6, y1),
                      color, -1)
        cv2.putText(output, text,
                    (x1 + 3, y1 - 4),
                    font, font_scale, (0, 0, 0), thickness, cv2.LINE_AA)

    return output


def print_summary(detections):
    counts = Counter(d['label'] for d in detections)
    print("\n-- Detection Summary ---------------------")
    for label, count in sorted(counts.items(), key=lambda x: -x[1]):
        avg_conf = np.mean([d['confidence'] for d in detections
                            if d['label'] == label])
        print(f"   {label:<20} x{count}   avg conf: {avg_conf:.0%}")
    print(f"   {'TOTAL':<20} x{len(detections)}")
    print("------------------------------------------\n")


def save_json(detections, output_path):
    data = {
        "total_components": len(detections),
        "components": [
            {**d, "bbox": list(d["bbox"])}
            for d in detections
        ]
    }
    with open(output_path, "w") as f:
        json.dump(data, f, indent=2)
    print(f"[OK] Results saved as JSON: {output_path}")


def detect_components(image_path, save_output=True, show_traces=True):
    print(f"\n{'='*50}")
    print(f"  PCB Component Detector")
    print(f"  Image: {image_path}")
    print(f"{'='*50}\n")

    img = cv2.imread(image_path)
    if img is None:
        print(f"[X] Could not load image: {image_path}")
        return []
    print(f"[OK] Image loaded: {img.shape[1]}x{img.shape[0]} px")

    detections = run_detection(image_path)

    if not detections:
        print("[!] No components detected. Try a clearer PCB image.")
        return []

    trace_mask = None
    if show_traces:
        trace_mask = detect_traces(img)
        trace_px   = np.count_nonzero(trace_mask)
        print(f"[OK] Traces detected: {trace_px} pixels")

    print_summary(detections)

    if save_output:
        base, ext = os.path.splitext(image_path)

        annotated = draw_detections(img, detections, trace_mask)
        img_out   = base + "_detected" + ext
        cv2.imwrite(img_out, annotated)
        print(f"[OK] Annotated image saved: {img_out}")

        json_out = base + "_results.json"
        save_json(detections, json_out)

    return detections


if __name__ == "__main__":
    if len(sys.argv) < 2:
        test_image = "sample 5.jpg"
        print(f"No image specified. Using default: {test_image}")
    else:
        test_image = sys.argv[1]

    detections = detect_components(test_image)