import gradio as gr
from huggingface_hub import snapshot_download
import sys

# =========================
# Load model + inspector
# =========================
MODEL_REPO = "Janani-V/pcb-defect-yolov8s-deeppcb"

local_dir = snapshot_download(repo_id=MODEL_REPO)
sys.path.append(local_dir)

from inspector import PCBDefectInspector

inspector = PCBDefectInspector(weights_path=f"{local_dir}/best.pt")

# =========================
# Severity config
# =========================
SEVERITY_COLORS = {
    "Critical": "#ef4444",
    "High": "#f97316",
    "Medium-High": "#f59e0b",
    "Medium": "#eab308",
    "Low-Medium": "#84cc16",
    "Low": "#22c55e",
    "Unknown": "#64748b",
}

SEVERITY_RANK = {
    "Critical": 6,
    "High": 5,
    "Medium-High": 4,
    "Medium": 3,
    "Low-Medium": 2,
    "Low": 1,
    "Unknown": 0,
}

# =========================
# Example images
# Make sure these exist in your Space repo
# =========================
EXAMPLES = [
    "sample1_input.jpg",
    "sample2_input.jpg",
    "sample3_input.jpg",
    "sample4_input.jpg",
    "sample5_input.jpg",
    "sample6_input.jpg",
    "sample7_input.jpg",
    "sample8_input.jpg",
]

# =========================
# Helper functions
# =========================
def empty_summary_cards():
    return """
    <div class="summary-grid">
        <div class="summary-card">
            <div class="summary-title">Defects Found</div>
            <div class="summary-value">-</div>
        </div>
        <div class="summary-card">
            <div class="summary-title">Highest Severity</div>
            <div class="summary-value">-</div>
        </div>
        <div class="summary-card">
            <div class="summary-title">Unique Defect Types</div>
            <div class="summary-value">-</div>
        </div>
        <div class="summary-card">
            <div class="summary-title">Avg Confidence</div>
            <div class="summary-value">-</div>
        </div>
    </div>
    """

def empty_report():
    return """
    <div class="empty-report">
        <h3>Inspection report will appear here</h3>
        <p>Upload a PCB image and run inspection to view defect analysis, severity, likely root cause, impact, and recommended action.</p>
    </div>
    """

def make_summary_cards(findings):
    if not findings:
        return """
        <div class="summary-grid">
            <div class="summary-card">
                <div class="summary-title">Defects Found</div>
                <div class="summary-value">0</div>
            </div>
            <div class="summary-card">
                <div class="summary-title">Highest Severity</div>
                <div class="summary-value">None</div>
            </div>
            <div class="summary-card">
                <div class="summary-title">Unique Defect Types</div>
                <div class="summary-value">0</div>
            </div>
            <div class="summary-card">
                <div class="summary-title">Avg Confidence</div>
                <div class="summary-value">0.00</div>
            </div>
        </div>
        """

    total_defects = len(findings)
    unique_classes = len(set(f["class"] for f in findings))
    avg_conf = sum(float(f["confidence"]) for f in findings) / max(total_defects, 1)

    highest_severity = max(
        (f.get("severity", "Unknown") for f in findings),
        key=lambda s: SEVERITY_RANK.get(s, 0),
        default="Unknown"
    )
    sev_color = SEVERITY_COLORS.get(highest_severity, "#64748b")

    return f"""
    <div class="summary-grid">
        <div class="summary-card">
            <div class="summary-title">Defects Found</div>
            <div class="summary-value">{total_defects}</div>
        </div>
        <div class="summary-card">
            <div class="summary-title">Highest Severity</div>
            <div class="summary-value" style="color:{sev_color};">{highest_severity}</div>
        </div>
        <div class="summary-card">
            <div class="summary-title">Unique Defect Types</div>
            <div class="summary-value">{unique_classes}</div>
        </div>
        <div class="summary-card">
            <div class="summary-title">Avg Confidence</div>
            <div class="summary-value">{avg_conf:.2f}</div>
        </div>
    </div>
    """

def make_report_html(result):
    findings = result.get("findings", [])
    summary = result.get("summary", "Inspection completed.")

    if not findings:
        return f"""
        <div class="report-wrapper">
            <div class="report-banner neutral">
                <div class="report-banner-title">Inspection Summary</div>
                <div class="report-banner-text">{summary}</div>
            </div>

            <div class="empty-report">
                <h3>No defects detected above the selected confidence threshold.</h3>
                <p>The model did not identify any PCB manufacturing defect in the uploaded image.</p>
            </div>
        </div>
        """

    # Group by class so repeated instances are summarized together
    grouped = {}
    for f in findings:
        cls = f["class"]
        if cls not in grouped:
            grouped[cls] = {
                "count": 0,
                "confidences": [],
                "info": f
            }
        grouped[cls]["count"] += 1
        grouped[cls]["confidences"].append(f["confidence"])

    sorted_items = sorted(
        grouped.items(),
        key=lambda x: SEVERITY_RANK.get(x[1]["info"].get("severity", "Unknown"), 0),
        reverse=True
    )

    highest_severity = sorted_items[0][1]["info"].get("severity", "Unknown")
    banner_color = SEVERITY_COLORS.get(highest_severity, "#64748b")

    html = f"""
    <div class="report-wrapper">
        <div class="report-banner" style="border-left: 6px solid {banner_color};">
            <div class="report-banner-title">Inspection Summary</div>
            <div class="report-banner-text">{summary}</div>
            <div class="report-banner-sub">Highest severity detected: <b>{highest_severity}</b></div>
        </div>
    """

    for cls, data in sorted_items:
        info = data["info"]
        severity = info.get("severity", "Unknown")
        sev_color = SEVERITY_COLORS.get(severity, "#64748b")

        conf_list = ", ".join(f"{float(c):.2f}" for c in data["confidences"])

        explanation = info.get("explanation", "No explanation available.")
        root_cause = info.get("root_cause", "Not available.")
        impact = info.get("impact", "Not available.")
        action = info.get("action", "Not available.")

        html += f"""
        <div class="defect-card">
            <div class="defect-header">
                <div>
                    <div class="defect-name">{cls.upper()}</div>
                    <div class="defect-count">Detected instances: {data['count']}</div>
                </div>
                <div class="severity-pill" style="background:{sev_color}; color:#111827;">
                    {severity}
                </div>
            </div>

            <div class="defect-meta">
                <span><b>Confidence:</b> {conf_list}</span>
            </div>

            <div class="defect-section">
                <div class="defect-section-title">Explanation</div>
                <div class="defect-section-text">{explanation}</div>
            </div>

            <div class="defect-section">
                <div class="defect-section-title">Likely Root Cause</div>
                <div class="defect-section-text">{root_cause}</div>
            </div>

            <div class="defect-section">
                <div class="defect-section-title">Potential Impact</div>
                <div class="defect-section-text">{impact}</div>
            </div>

            <div class="defect-section">
                <div class="defect-section-title">Recommended Action</div>
                <div class="defect-section-text">{action}</div>
            </div>
        </div>
        """

    html += "</div>"
    return html


# =========================
# Main functions
# =========================
def detect(image, conf):
    if image is None:
        return None, empty_summary_cards(), empty_report()

    inspector.conf = conf
    result = inspector.inspect(image)

    annotated = result.get("annotated_image", None)
    findings = result.get("findings", [])

    summary_cards = make_summary_cards(findings)
    report_html = make_report_html(result)

    return annotated, summary_cards, report_html


def clear_all():
    return None, 0.25, None, empty_summary_cards(), empty_report()


# =========================
# CSS
# =========================
CUSTOM_CSS = """
body, .gradio-container {
    background: #0b1220 !important;
    color: #f3f4f6 !important;
    font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
}

.gradio-container {
    max-width: 1400px !important;
    margin: auto !important;
    padding-top: 18px !important;
    padding-bottom: 24px !important;
}

.main-title {
    font-size: 34px;
    font-weight: 800;
    margin-bottom: 6px;
    color: #f8fafc;
}

.sub-title {
    font-size: 16px;
    line-height: 1.8;
    color: #cbd5e1;
    margin-bottom: 24px;
}

.section-title {
    font-size: 22px;
    font-weight: 700;
    color: #f8fafc;
    margin: 8px 0 14px 0;
}

.card-subtext {
    color: #94a3b8;
    font-size: 14px;
    margin-top: 8px;
    line-height: 1.7;
}

.summary-grid {
    display: grid;
    grid-template-columns: repeat(4, minmax(120px, 1fr));
    gap: 12px;
    margin-top: 12px;
    margin-bottom: 6px;
}

.summary-card {
    background: linear-gradient(180deg, #141d2d 0%, #0f172a 100%);
    border: 1px solid #243042;
    border-radius: 16px;
    padding: 16px;
    min-height: 92px;
    display: flex;
    flex-direction: column;
    justify-content: center;
}

.summary-title {
    color: #94a3b8;
    font-size: 13px;
    margin-bottom: 8px;
}

.summary-value {
    font-size: 24px;
    font-weight: 800;
    color: #f8fafc;
}

.report-wrapper {
    margin-top: 6px;
}

.report-banner {
    background: #111827;
    border: 1px solid #243042;
    border-radius: 16px;
    padding: 18px 18px 16px 18px;
    margin-bottom: 16px;
}

.report-banner-title {
    font-size: 18px;
    font-weight: 800;
    color: #f8fafc;
    margin-bottom: 6px;
}

.report-banner-text {
    color: #dbeafe;
    font-size: 15px;
    line-height: 1.7;
    margin-bottom: 6px;
}

.report-banner-sub {
    color: #cbd5e1;
    font-size: 14px;
}

.defect-card {
    background: #111827;
    border: 1px solid #243042;
    border-radius: 16px;
    padding: 18px;
    margin-bottom: 14px;
}

.defect-header {
    display: flex;
    justify-content: space-between;
    align-items: flex-start;
    gap: 12px;
    margin-bottom: 12px;
}

.defect-name {
    font-size: 20px;
    font-weight: 800;
    color: #f8fafc;
}

.defect-count {
    color: #94a3b8;
    font-size: 14px;
    margin-top: 4px;
}

.severity-pill {
    padding: 8px 14px;
    border-radius: 999px;
    font-weight: 800;
    font-size: 13px;
    white-space: nowrap;
}

.defect-meta {
    font-size: 14px;
    color: #cbd5e1;
    margin-bottom: 14px;
}

.defect-section {
    margin-top: 12px;
    padding-top: 10px;
    border-top: 1px solid #223045;
}

.defect-section-title {
    font-size: 14px;
    font-weight: 800;
    color: #93c5fd;
    margin-bottom: 6px;
    letter-spacing: 0.2px;
}

.defect-section-text {
    color: #e5e7eb;
    line-height: 1.7;
    font-size: 14px;
}

.empty-report {
    background: #111827;
    border: 1px dashed #334155;
    border-radius: 16px;
    padding: 28px;
    text-align: center;
    color: #cbd5e1;
    margin-top: 8px;
}

.empty-report h3 {
    margin: 0 0 10px 0;
    color: #f8fafc;
    font-size: 20px;
}

.note-box {
    background: #0f172a;
    border: 1px solid #223045;
    border-radius: 16px;
    padding: 16px;
    color: #cbd5e1;
    font-size: 14px;
    line-height: 1.7;
    margin-top: 18px;
}

.note-box b {
    color: #f8fafc;
}

.footer-link {
    color: #8b5cf6 !important;
    text-decoration: none !important;
    font-weight: 700;
}

.footer-link:hover {
    text-decoration: underline !important;
}

/* Make buttons cleaner */
button.primary {
    background: linear-gradient(90deg, #6D5DF6, #7C3AED) !important;
    border: none !important;
    color: white !important;
}

button.secondary {
    background: #1f2937 !important;
    border: 1px solid #374151 !important;
    color: #f8fafc !important;
}

@media (max-width: 900px) {
    .summary-grid {
        grid-template-columns: repeat(2, minmax(120px, 1fr));
    }
}

@media (max-width: 560px) {
    .summary-grid {
        grid-template-columns: 1fr;
    }
}

/* Even, uniform example thumbnails (Gradio 5 renders Examples as a
   [data-testid="dataset"] grid, not a <table>) */
.gradio-container [data-testid="dataset"] {
    display: grid !important;
    grid-template-columns: repeat(4, minmax(90px, 1fr)) !important;
    gap: 14px !important;
}

.gradio-container [data-testid="dataset"] > div,
.gradio-container [data-testid="dataset"] button {
    width: 100% !important;
    aspect-ratio: 1 / 1 !important;
    padding: 0 !important;
    margin: 0 !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    overflow: hidden !important;
    border-radius: 10px !important;
}

.gradio-container [data-testid="dataset"] img {
    width: 100% !important;
    height: 100% !important;
    object-fit: cover !important;
    border-radius: 10px !important;
}

@media (max-width: 900px) {
    .gradio-container [data-testid="dataset"] {
        grid-template-columns: repeat(2, minmax(90px, 1fr)) !important;
    }
}
"""

# =========================
# Build UI
# =========================
with gr.Blocks(css=CUSTOM_CSS) as demo:
    gr.HTML("""
    <div class="main-title">PCB Defect Detection Dashboard</div>
    <div class="sub-title">
        Upload a PCB image to detect manufacturing defects and receive an inspection report with
        <b>severity assessment</b>, <b>likely root cause</b>, <b>impact</b>, and
        <b>recommended action</b> for each detected defect type.
        This interface combines a fine-tuned <b>YOLOv8s defect detector</b> with a structured PCB defect interpretation layer.
    </div>
    """)

    with gr.Row(equal_height=False):
        with gr.Column(scale=4):
            gr.HTML('<div class="section-title">📤 Input Panel</div>')

            input_image = gr.Image(
                type="numpy",
                label="Upload PCB Image",
                height=360
            )

            conf_slider = gr.Slider(
                minimum=0.10,
                maximum=0.90,
                value=0.25,
                step=0.05,
                label="Confidence Threshold"
            )

            with gr.Row():
                detect_btn = gr.Button("Run Inspection", variant="primary")
                clear_btn = gr.Button("Clear", variant="secondary")

            gr.HTML("""
            <div class="card-subtext">
                Supported formats: <b>JPG, JPEG, PNG</b><br>
                Recommended: clear grayscale PCB images similar to DeepPCB-style inspection images.
            </div>
            """)

        # RIGHT PANEL
        with gr.Column(scale=6):
            gr.HTML('<div class="section-title">🔍 Detection Output</div>')

            output_image = gr.Image(
                type="numpy",
                label="Annotated PCB Output",
                height=360
            )

            summary_html = gr.HTML(empty_summary_cards())

    gr.HTML('<div class="section-title" style="margin-top:18px;">📋 Intelligent Inspection Report</div>')
    report_html = gr.HTML(empty_report())

    gr.HTML('<div class="section-title" style="margin-top:20px;">🖼 Example PCB Images</div>')
    gr.Examples(
        examples=EXAMPLES,
        inputs=input_image,
        label="Click any example to load it into the input panel"
    )

    gr.HTML(f"""
    <div class="note-box">
        <b>Note:</b> Severity, root cause, impact, and recommended action are generated from a structured PCB defect knowledge layer built on top of the model's class predictions.
        The YOLO model itself performs visual defect detection; higher-level interpretation should still be reviewed in the context of PCB manufacturing and inspection workflows.
        <br><br>
        <b>Model card:</b>
        <a class="footer-link" href="https://huggingface.co/{MODEL_REPO}" target="_blank">
            {MODEL_REPO}
        </a>
    </div>
    """)

    detect_btn.click(
        fn=detect,
        inputs=[input_image, conf_slider],
        outputs=[output_image, summary_html, report_html]
    )

    clear_btn.click(
        fn=clear_all,
        inputs=[],
        outputs=[input_image, conf_slider, output_image, summary_html, report_html]
    )

if __name__ == "__main__":
    demo.launch()