import gradio as gr
import cv2
import numpy as np
import json
import os
import tempfile
import shutil
from PIL import Image
from collections import Counter
from inference_sdk import InferenceHTTPClient
import easyocr

# ─────────────────────────────────────────────
#  IMPORTS FROM OUR PIPELINE
# ─────────────────────────────────────────────
from detector import run_detection, detect_traces, draw_detections
from ocr import run_ocr_on_detections, print_ocr_summary
from netlist import (assign_reference_designators,
                     extract_trace_mask,
                     find_trace_connections,
                     find_proximity_connections,
                     build_nets)
from kicad_writer import generate_kicad_schematic

# ─────────────────────────────────────────────
#  GLOBAL OCR READER (load once)
# ─────────────────────────────────────────────
print("[->] Loading EasyOCR...")
ocr_reader = easyocr.Reader(['en'], gpu=False)  # CPU for HuggingFace
print("[OK] EasyOCR ready")


# ─────────────────────────────────────────────
#  HELPER — numpy image to PIL
# ─────────────────────────────────────────────
def to_pil(img_bgr: np.ndarray) -> Image.Image:
    return Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))


# ─────────────────────────────────────────────
#  TAB 1 — COMPONENT DETECTION
# ─────────────────────────────────────────────
def run_detection_tab(image: Image.Image):
    if image is None:
        return None, "❌ Please upload a PCB image first.", "{}"

    # Save uploaded image to temp file
    tmp_dir  = tempfile.mkdtemp()
    img_path = os.path.join(tmp_dir, "input.jpg")
    image.save(img_path)

    try:
        # Run Roboflow detection
        detections = run_detection(img_path)

        if not detections:
            return image, "⚠️ No components detected. Try a clearer PCB image.", "{}"

        # Draw detections
        img_bgr   = cv2.imread(img_path)
        annotated = draw_detections(img_bgr, detections)
        result_pil= to_pil(annotated)

        # Build summary text
        counts   = Counter(d['label'] for d in detections)
        summary  = f"✅ **{len(detections)} components detected**\n\n"
        summary += "| Component | Count |\n|-----------|-------|\n"
        for label, count in sorted(counts.items(), key=lambda x: -x[1]):
            summary += f"| {label} | {count} |\n"

        # Save detections to JSON for next tabs
        det_json = json.dumps({
            "image_path": img_path,
            "components": [
                {**d, "bbox": list(d["bbox"])}
                for d in detections
            ]
        }, indent=2)

        return result_pil, summary, det_json

    except Exception as e:
        return image, f"❌ Error: {str(e)}", "{}"


# ─────────────────────────────────────────────
#  TAB 2 — OCR / PART NUMBER READING
# ─────────────────────────────────────────────
def run_ocr_tab(det_json: str):
    if not det_json or det_json == "{}":
        return "⚠️ Run Detection first!", "{}"

    try:
        data       = json.loads(det_json)
        img_path   = data.get("image_path")
        detections = data.get("components", [])

        for d in detections:
            d['bbox'] = tuple(d['bbox'])

        IC_LABELS = ['ic', 'transistor', 'clock', 'display']

        # Run OCR using global reader
        img = cv2.imread(img_path)
        ih, iw = img.shape[:2]
        PADDING = 10
        updated = []

        for det in detections:
            label = det['label']
            if label not in IC_LABELS:
                det['ocr_text']    = []
                det['part_number'] = "N/A"
                updated.append(det)
                continue

            x1, y1, x2, y2 = det['bbox']
            x1p = max(0,  x1 - PADDING)
            y1p = max(0,  y1 - PADDING)
            x2p = min(iw, x2 + PADDING)
            y2p = min(ih, y2 + PADDING)
            patch = img[y1p:y2p, x1p:x2p]

            if patch.size == 0:
                det['ocr_text']    = []
                det['part_number'] = "unknown"
                updated.append(det)
                continue

            # Upscale for better OCR
            h, w   = patch.shape[:2]
            scale  = 3 if max(h, w) < 100 else 2
            patch  = cv2.resize(patch, (w*scale, h*scale),
                                interpolation=cv2.INTER_CUBIC)

            results  = ocr_reader.readtext(patch)
            texts    = [(t.strip(), round(c, 3))
                        for _, t, c in results
                        if c >= 0.4 and len(t.strip()) >= 2]

            combined           = " ".join(t for t, c in texts).strip()
            det['ocr_text']    = texts
            det['part_number'] = combined if combined else "unknown"
            updated.append(det)

        # Build output table
        ic_dets    = [d for d in updated if d['label'] in IC_LABELS]
        identified = [d for d in ic_dets
                      if d.get('part_number', 'unknown') not in
                      ('unknown', 'N/A', '')]

        summary  = f"✅ **OCR complete — {len(identified)}/{len(ic_dets)} ICs identified**\n\n"
        summary += "| RefDes | Label | Part Number | Confidence |\n"
        summary += "|--------|-------|-------------|------------|\n"

        for i, det in enumerate(ic_dets):
            ref  = f"U{i+1}"
            part = det.get('part_number', 'unknown')
            conf = det['confidence']
            summary += f"| {ref} | {det['label']} | {part} | {conf:.0%} |\n"

        # Pass updated detections forward
        out_json = json.dumps({
            "image_path": img_path,
            "components": [
                {**d,
                 "bbox":     list(d["bbox"]),
                 "ocr_text": [[t, c] for t, c in d.get("ocr_text", [])]}
                for d in updated
            ]
        }, indent=2)

        return summary, out_json

    except Exception as e:
        return f"❌ Error: {str(e)}", "{}"


# ─────────────────────────────────────────────
#  TAB 3 — NETLIST GENERATION
# ─────────────────────────────────────────────
def run_netlist_tab(ocr_json: str):
    if not ocr_json or ocr_json == "{}":
        return "⚠️ Run OCR first!", "{}", None

    try:
        data       = json.loads(ocr_json)
        img_path   = data.get("image_path")
        detections = data.get("components", [])

        for d in detections:
            d['bbox'] = tuple(d['bbox'])

        img        = cv2.imread(img_path)
        components = assign_reference_designators(detections)
        trace_mask = extract_trace_mask(img)
        trace_conn = find_trace_connections(components, trace_mask, img.shape)
        prox_conn  = find_proximity_connections(components, trace_conn)
        all_conn   = trace_conn + prox_conn
        nets       = build_nets(all_conn)

        # Save netlist JSON to temp file
        tmp_dir      = os.path.dirname(img_path)
        netlist_path = os.path.join(tmp_dir, "netlist.json")

        out_data = {
            "total_components": len(components),
            "total_connections": len(all_conn),
            "total_nets": len(nets),
            "components": [
                {**c,
                 "bbox":     list(c["bbox"]),
                 "ocr_text": [[t, conf] for t, conf
                              in c.get("ocr_text", [])]}
                for c in components
            ],
            "connections": [
                {"from": a, "to": b, "method": m}
                for a, b, m in all_conn
            ],
            "nets": nets
        }

        with open(netlist_path, "w") as f:
            json.dump(out_data, f, indent=2)

        # Build summary
        trace_c = len([c for c in all_conn if c[2] == "trace"])
        prox_c  = len([c for c in all_conn if c[2] == "proximity"])

        summary  = f"✅ **Netlist generated successfully**\n\n"
        summary += f"- **Components:** {len(components)}\n"
        summary += f"- **Connections:** {len(all_conn)} "
        summary += f"({trace_c} via traces, {prox_c} via proximity)\n"
        summary += f"- **Nets:** {len(nets)}\n\n"
        summary += "| Net | Members |\n|-----|--------|\n"
        for net_name, members in nets.items():
            summary += f"| {net_name} | {', '.join(members[:5])}"
            if len(members) > 5:
                summary += f" ... +{len(members)-5} more"
            summary += " |\n"

        return summary, json.dumps({"netlist_path": netlist_path}), netlist_path

    except Exception as e:
        return f"❌ Error: {str(e)}", "{}", None


# ─────────────────────────────────────────────
#  TAB 4 — KICAD SCHEMATIC OUTPUT
# ─────────────────────────────────────────────
def run_kicad_tab(netlist_ref: str):
    if not netlist_ref or netlist_ref == "{}":
        return "⚠️ Run Netlist generation first!", None

    try:
        data         = json.loads(netlist_ref)
        netlist_path = data.get("netlist_path")

        if not netlist_path or not os.path.exists(netlist_path):
            return "❌ Netlist file not found. Re-run previous steps.", None

        # Generate KiCAD schematic
        tmp_dir    = os.path.dirname(netlist_path)
        sch_path   = os.path.join(tmp_dir, "schematic.kicad_sch")

        generate_kicad_schematic(netlist_path, sch_path)

        summary  = f"✅ **KiCAD schematic generated!**\n\n"
        summary += f"- File: `schematic.kicad_sch`\n"
        summary += f"- Format: KiCAD 6/7 compatible\n\n"
        summary += "**How to open:**\n"
        summary += "1. Download the file below\n"
        summary += "2. Open KiCAD → File → Open Schematic\n"
        summary += "   OR drag into [kicanvas.org](https://kicanvas.org) for instant preview\n"

        return summary, sch_path

    except Exception as e:
        return f"❌ Error: {str(e)}", None


# ─────────────────────────────────────────────
#  BUILD GRADIO UI
# ─────────────────────────────────────────────
def build_ui():
    with gr.Blocks(
        title="PCB Image → Schematic",
        theme=gr.themes.Soft(),
        css="""
        .tab-header { font-size: 1.1em; font-weight: bold; }
        .output-panel { background: #1a1a2e; border-radius: 8px; }
        """
    ) as demo:

        # ── Header ──
        gr.Markdown("""
        # 🔌 PCB Image → Schematic
        ### Convert a PCB photo into a KiCAD schematic automatically
        Upload a PCB image and step through each stage of the pipeline.
        """)

        # ── Shared state between tabs ──
        detection_state = gr.State("{}")
        ocr_state       = gr.State("{}")
        netlist_state   = gr.State("{}")

        # ── Tab 1: Detection ──
        with gr.Tab("📷 1 — Component Detection"):
            gr.Markdown("Upload a PCB image or click one of the example images below.")
            with gr.Row():
                with gr.Column(scale=1):
                    img_input  = gr.Image(type="pil", label="PCB Image")
                    detect_btn = gr.Button("🔍 Detect Components", variant="primary")
                    gr.Examples(
                        examples=[
                            ["sample 1.jpg"],
                            ["sample 2.jpg"],
                            ["sample 3.jpg"],
                            ["sample 4.jpg"],
                            ["sample 5.jpg"],
                        ],
                        inputs=img_input,
                        label="📂 Example PCB Images — click to load",
                        examples_per_page=5,
                    )
                with gr.Column(scale=1):
                    detect_out  = gr.Image(label="Detected Components")
                    detect_text = gr.Markdown()

            detect_btn.click(
                fn=run_detection_tab,
                inputs=[img_input],
                outputs=[detect_out, detect_text, detection_state]
            )

        # ── Tab 2: OCR ──
        with gr.Tab("🔤 2 — Read IC Text"):
            gr.Markdown("Reads part numbers from IC chips using OCR.")
            ocr_btn  = gr.Button("📖 Run OCR on ICs", variant="primary")
            ocr_text = gr.Markdown()

            ocr_btn.click(
                fn=run_ocr_tab,
                inputs=[detection_state],
                outputs=[ocr_text, ocr_state]
            )

        # ── Tab 3: Netlist ──
        with gr.Tab("🔗 3 — Generate Netlist"):
            gr.Markdown("Finds connections between components using trace detection + proximity.")
            netlist_btn  = gr.Button("⚡ Generate Netlist", variant="primary")
            netlist_text = gr.Markdown()
            netlist_file = gr.File(label="Download Netlist JSON", visible=False)

            netlist_btn.click(
                fn=run_netlist_tab,
                inputs=[ocr_state],
                outputs=[netlist_text, netlist_state, netlist_file]
            )

        # ── Tab 4: KiCAD ──
        with gr.Tab("📐 4 — KiCAD Schematic"):
            gr.Markdown("Generates a KiCAD `.kicad_sch` file you can open in KiCAD or kicanvas.org")
            kicad_btn  = gr.Button("💾 Generate KiCAD File", variant="primary")
            kicad_text = gr.Markdown()
            kicad_file = gr.File(label="Download .kicad_sch")

            kicad_btn.click(
                fn=run_kicad_tab,
                inputs=[netlist_state],
                outputs=[kicad_text, kicad_file]
            )

        # ── Footer ──
        gr.Markdown("""
        ---
        Built with Roboflow YOLOv8 · EasyOCR · OpenCV · KiCAD
        """)

    return demo


# ─────────────────────────────────────────────
#  ENTRY POINT
# ─────────────────────────────────────────────
if __name__ == "__main__":
    demo = build_ui()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )