"""
病虫害以图搜图
基于图片 Embedding API 的相似度搜索
"""

from __future__ import annotations

import io
import time
from dataclasses import dataclass
from typing import Literal

import httpx
import numpy as np
import plotly.graph_objects as go
import requests
import streamlit as st
from PIL import Image

# ─── API Config ──────────────────────────────────────────────────────────────
IMAGE_EMBEDDING_API_URL = "https://llm.dev.maimaiag.com/qwen3-vl-embedding/v1/embeddings"
EMBEDDING_MODEL = "Qwen3-VL-Embedding"
API_KEY = "sk--VnOesEU5D8wnHjdg0MEsA"

# ─── Page Config ────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="病虫害以图搜图",
    page_icon="🌿",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ─── Minimal CSS ─────────────────────────────────────────────────────────────
st.markdown("""
<style>
html, body, [class*="css"] {
    font-family: "PingFang SC", "Microsoft YaHei", "Noto Sans SC", sans-serif;
}
</style>
""", unsafe_allow_html=True)


# ─── Knowledge Base ──────────────────────────────────────────────────────────
@dataclass(frozen=True)
class PestItem:
    name: str
    url: str
    symptoms: str
    treatment: str
    crop: str
    category: Literal["病害", "虫害"]


PEST_KNOWLEDGE: list[PestItem] = [
    PestItem(
        name="水稻稻瘟病",
        url="https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_151854_dc9667cf_%E6%B0%B4%E7%A8%BB%E7%A8%BB%E7%98%9F%E7%97%851.jpeg",
        symptoms="叶片出现梭形或纺锤形病斑，中央灰白色，边缘褐色，严重时病斑连片导致叶片枯死",
        treatment="选用抗病品种，合理施肥避免偏施氮肥，发病初期喷施三环唑或稻瘟灵",
        crop="水稻",
        category="病害",
    ),
    PestItem(
        name="水稻纹枯病",
        url="https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_152022_9f3124ab_%E6%B0%B4%E7%A8%BB%E7%BA%B9%E6%9E%AF%E7%97%851.jpeg",
        symptoms="叶鞘和叶片上出现云纹状灰绿色至灰褐色病斑，后期病斑边缘褐色、中央灰白色",
        treatment="合理密植，科学管水，发病初期喷施井冈霉素或噻呋酰胺",
        crop="水稻",
        category="病害",
    ),
    PestItem(
        name="水稻胡麻叶斑病",
        url="https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_151936_41fdb1dc_%E6%B0%B4%E7%A8%BB%E8%83%A1%E9%BA%BB%E5%8F%B6%E6%96%91%E7%97%851.jpeg",
        symptoms="叶片上出现暗褐色芝麻粒大小的椭圆形病斑，病斑周围有黄色晕圈",
        treatment="增施硅肥和钾肥提高抗病力，喷施丙环唑或咪鲜胺防治",
        crop="水稻",
        category="病害",
    ),
    PestItem(
        name="小麦锈病",
        url="https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_153814_3e175ca3_%E5%B0%8F%E9%BA%A6%E9%94%88%E7%97%851.jpeg",
        symptoms="叶片和叶鞘上出现铁锈色粉状疱疹（夏孢子堆），后期变为黑色冬孢子堆",
        treatment="种植抗锈品种，发病初期喷施三唑酮或烯唑醇，注意轮作",
        crop="小麦",
        category="病害",
    ),
    PestItem(
        name="小麦赤霉病",
        url="https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_152112_2e1f530e_%E5%B0%8F%E9%BA%A6%E8%B5%A4%E9%9C%89%E7%97%851.jpeg",
        symptoms="穗部小穗发病，颖壳上出现水浸状褐色斑，后期产生粉红色霉层",
        treatment="选用抗病品种，齐穗至扬花初期喷施多菌灵或戊唑醇",
        crop="小麦",
        category="病害",
    ),
    PestItem(
        name="玉米大斑病",
        url="https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_153911_ee5a72be_%E7%8E%89%E7%B1%B3%E5%A4%A7%E6%96%91%E7%97%851.jpeg",
        symptoms="叶片上出现灰绿色水浸状斑点，扩展为长梭形灰褐色大型病斑",
        treatment="种植抗病品种，适时早播，发病初期喷施多菌灵或代森锰锌",
        crop="玉米",
        category="病害",
    ),
    PestItem(
        name="玉米小斑病",
        url="https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_154001_e31a0103_%E7%8E%89%E7%B1%B3%E5%B0%8F%E6%96%91%E7%97%851.jpeg",
        symptoms="叶片上出现椭圆形黄褐色小病斑，有2-3圈同心轮纹，边缘紫褐色",
        treatment="轮作倒茬，清除病残体，喷施百菌清或甲基托布津",
        crop="玉米",
        category="病害",
    ),
    PestItem(
        name="玉米螟",
        url="https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_153938_8be05006_%E7%8E%89%E7%B1%B3%E8%9E%9F1.jpeg",
        symptoms="幼虫蛀食茎秆和穗轴，茎秆上有蛀孔，孔口有虫粪，造成茎秆折断",
        treatment="心叶期撒施白僵菌颗粒剂，释放赤眼蜂生物防治，大喇叭口期灌心",
        crop="玉米",
        category="虫害",
    ),
    PestItem(
        name="稻飞虱",
        url="https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_151643_db5e1d36_%E7%A8%BB%E9%A3%9E%E8%99%AB1.jpeg",
        symptoms="稻株基部聚集大量褐色或白色小型飞虫，受害稻株发黄矮缩，严重时枯死倒伏",
        treatment="合理施肥避免贪青晚熟，选用吡蚜酮或烯啶虫胺防治，保护利用天敌",
        crop="水稻",
        category="虫害",
    ),
    PestItem(
        name="大豆蚜虫",
        url="https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_151549_d9cf327b_%E5%A4%A7%E8%B1%86%E8%9A%9C%E8%99%AB1.jpeg",
        symptoms="嫩叶和茎尖聚集大量绿色或黄色蚜虫，叶片卷缩变形，植株矮化",
        treatment="保护瓢虫等天敌，百株蚜量达1000头时喷施吡虫啉或啶虫脒",
        crop="大豆",
        category="虫害",
    ),
    PestItem(
        name="番茄晚疫病",
        url="https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_151705_3dd8baab_%E7%95%AA%E8%8C%84%E6%99%9A%E7%96%AB%E7%97%851.jpeg",
        symptoms="叶片出现水浸状暗绿色不规则病斑，潮湿时叶背面产生白色霉层，果实变褐硬化",
        treatment="控制温湿度，及时通风降湿，发病初期喷施甲霜灵锰锌或霜脲氰",
        crop="番茄",
        category="病害",
    ),
    PestItem(
        name="黄瓜霜霉病",
        url="https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_151804_7be515fa_%E9%BB%84%E7%93%9C%E9%9C%9C%E9%9C%89%E7%97%851.jpeg",
        symptoms="叶片正面出现黄色多角形病斑，叶背面潮湿时产生灰黑色霉层",
        treatment="选用抗病品种，膜下滴灌降低湿度，喷施百菌清或霜霉威盐酸盐",
        crop="黄瓜",
        category="病害",
    ),
]

EXAMPLE_IMAGES: list[tuple[str, str]] = [
    (
        "水稻稻瘟病",
        "https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_151914_4f5b8fef_%E6%B0%B4%E7%A8%BB%E7%A8%BB%E7%98%9F%E7%97%852.jpeg",
    ),
    (
        "番茄晚疫病",
        "https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_151726_a8f31320_%E7%95%AA%E8%8C%84%E6%99%9A%E7%96%AB%E7%97%852.jpeg",
    ),
    (
        "小麦锈病",
        "https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_153837_e8ae9f43_%E5%B0%8F%E9%BA%A6%E9%94%88%E7%97%852.jpeg",
    ),
    (
        "水稻纹枯病",
        "https://minio.dev.maimaiag.com/crop-prod-bucket/field_photo/20260410_152050_77d568b1_%E6%B0%B4%E7%A8%BB%E7%BA%B9%E6%9E%AF%E7%97%852.jpeg",
    ),
]


# ─── Embedding API ───────────────────────────────────────────────────────────
def get_image_embedding(image_url: str, text: str = "这是什么病虫害？", max_retries: int = 2) -> list[float]:
    """调用远程 API 获取图片 Embedding，支持重试。"""
    payload = {
        "model": EMBEDDING_MODEL,
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": text},
                    {"type": "image_url", "image_url": {"url": image_url}},
                ],
            }
        ],
    }
    last_error = None
    for attempt in range(max_retries + 1):
        try:
            resp = httpx.post(
                IMAGE_EMBEDDING_API_URL,
                headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"},
                json=payload,
                timeout=120,
            )
            resp.raise_for_status()
            return resp.json()["data"][0]["embedding"]
        except httpx.HTTPStatusError as e:
            last_error = e
            if attempt < max_retries:
                time.sleep(2 * (attempt + 1))
    raise last_error


# ─── Utilities ───────────────────────────────────────────────────────────────
def _load_image_raw(source: str | io.BytesIO) -> Image.Image | None:
    try:
        if isinstance(source, str):
            resp = requests.get(source, timeout=30)
            resp.raise_for_status()
            return Image.open(io.BytesIO(resp.content)).convert("RGB")
        return Image.open(source).convert("RGB")
    except Exception:
        return None


def load_image(source: str | io.BytesIO) -> Image.Image | None:
    img = _load_image_raw(source)
    if img is None:
        st.error("图片加载失败，请检查链接是否可访问或文件是否损坏")
    return img


def cosine_similarity(a: list[float], b: list[float]) -> float:
    a_arr = np.array(a)
    b_arr = np.array(b)
    norm_a = np.linalg.norm(a_arr)
    norm_b = np.linalg.norm(b_arr)
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return float(np.dot(a_arr, b_arr) / (norm_a * norm_b))


@st.cache_resource
def build_index() -> tuple[list[dict], list[str], list[tuple[str, str]]]:
    items, succeeded, failed = [], [], []
    progress = st.progress(0, text="正在构建图片索引...")
    total = len(PEST_KNOWLEDGE)
    for i, pest in enumerate(PEST_KNOWLEDGE):
        try:
            embedding = get_image_embedding(pest.url, text=pest.name)
            items.append({
                "name": pest.name,
                "url": pest.url,
                "embedding": embedding,
                "symptoms": pest.symptoms,
                "treatment": pest.treatment,
                "crop": pest.crop,
                "category": pest.category,
            })
            succeeded.append(pest.name)
        except Exception as e:
            failed.append((pest.name, str(e)))
        progress.progress((i + 1) / total, text=f"正在构建图片索引 ({i + 1}/{total})...")
    progress.empty()
    return items, succeeded, failed


# ─── Sidebar ─────────────────────────────────────────────────────────────────
with st.sidebar:
    st.header("🌿 病虫害以图搜图")
    st.caption("输入图片 URL，通过图片 Embedding 搜索相似病虫害")
    st.divider()

    st.subheader("🖼️ 输入方式")
    input_mode = st.radio("", ["输入图片 URL", "选择示例图片"], label_visibility="collapsed")

    # 初始化 session_state
    if "query_url" not in st.session_state:
        st.session_state.query_url = ""

    query_source = None
    query_url = ""

    if input_mode == "输入图片 URL":
        query_url = st.text_input("图片 URL", value=st.session_state.query_url, placeholder="https://example.com/image.jpg")
        st.session_state.query_url = query_url
        if query_url.strip():
            query_source = query_url.strip()
    else:
        st.caption("点击选择示例")
        cols = st.columns(2)
        for idx, (name, url) in enumerate(EXAMPLE_IMAGES):
            with cols[idx % 2]:
                if st.button(name, key=f"ex_{name}"):
                    st.session_state.query_url = url
                    st.rerun()
        if st.session_state.query_url:
            query_url = st.session_state.query_url
            query_source = query_url
            st.image(query_url, use_container_width=True)

    st.subheader("⚙️ 搜索设置")
    top_k = st.slider("返回条数", 1, min(12, len(PEST_KNOWLEDGE)), 5)

    search_clicked = st.button("开始搜索", type="primary", use_container_width=True)

    st.divider()
    st.info(
        "**使用说明**\n\n"
        "1. 输入病虫害图片 URL 或选择示例\n"
        "2. 系统调用 Embedding API 提取图像特征\n"
        "3. 与知识库比对返回相似结果\n"
        "4. 参考症状与防治建议"
    )

# ─── Build Index ─────────────────────────────────────────────────────────────
with st.spinner("首次加载需要构建图片索引，请稍候..."):
    index_items, succeeded, failed = build_index()

# ─── Main Layout ─────────────────────────────────────────────────────────────
st.title("🌿 病虫害以图搜图")
st.caption("基于图片 Embedding API 的病虫害相似度检索与防治建议")

# Status badges
st.success(f"📚 图片索引构建完成，成功 {len(succeeded)} 张")
if failed:
    st.warning(f"以下 {len(failed)} 张图片索引失败：")
    for name, err in failed:
        st.error(f"- **{name}**：{err}")

st.markdown("<br>", unsafe_allow_html=True)

# ─── Search Logic ────────────────────────────────────────────────────────────
if search_clicked and query_url.strip():
    if not index_items:
        st.warning("知识库索引为空，请检查网络连接后刷新页面重试。")
    else:
        with st.spinner("正在分析图片并搜索..."):
            try:
                query_embedding = get_image_embedding(query_url.strip())

                col_query, col_preview = st.columns([1, 3])
                with col_query:
                    st.subheader("🔍 查询图片")
                    query_img = load_image(query_url.strip())
                    if query_img is not None:
                        st.image(query_img, use_container_width=True)

                with col_preview:
                    scores = []
                    for item in index_items:
                        sim = cosine_similarity(query_embedding, item["embedding"])
                        scores.append((sim, item))
                    scores.sort(key=lambda x: x[0], reverse=True)
                    results = scores[:top_k]

                    st.subheader(f"🏆 搜索结果（Top-{len(results)}）")

                    # Similarity bar chart
                    names = [f"{r[1]['name']}" for r in results]
                    sims = [r[0] * 100 for r in results]
                    colors = ["#c45c4a" if r[1]["category"] == "虫害" else "#4a7c59" for r in results]

                    fig_bar = go.Figure()
                    fig_bar.add_trace(go.Bar(
                        x=sims,
                        y=names,
                        orientation="h",
                        marker=dict(color=colors, opacity=0.85, line=dict(color="rgba(0,0,0,0.08)", width=1)),
                        text=[f"{s:.1f}%" for s in sims],
                        textposition="outside",
                        textfont=dict(color="#5a5a5a", size=10),
                    ))
                    fig_bar.update_layout(
                        xaxis=dict(title="相似度 (%)", color="#5a5a5a", gridcolor="rgba(0,0,0,0.06)", range=[0, 105]),
                        yaxis=dict(color="#5a5a5a", gridcolor="rgba(0,0,0,0.04)", autorange="reversed"),
                        paper_bgcolor="rgba(0,0,0,0)",
                        plot_bgcolor="rgba(0,0,0,0)",
                        font=dict(color="#2c2c2c", size=11),
                        margin=dict(t=10, b=30, l=80, r=50),
                        height=160 + len(results) * 34,
                        showlegend=False,
                    )
                    st.plotly_chart(fig_bar, use_container_width=True)

                # Result cards below
                st.subheader("📋 详细结果")
                for rank, (sim, item) in enumerate(results, 1):
                    with st.container(border=True):
                        c1, c2 = st.columns([1, 4])
                        with c1:
                            st.image(item["url"], use_container_width=True)
                        with c2:
                            header_col, score_col = st.columns([3, 1])
                            header_col.markdown(f"**#{rank} {item['name']}**")
                            score_col.markdown(f"<div style='text-align:right; font-weight:600;'>相似度 {sim*100:.1f}%</div>", unsafe_allow_html=True)

                            badge_cols = st.columns([1, 1, 4])
                            badge_cols[0].caption(f"🌾 {item['crop']}")
                            badge_cols[1].caption(f"🐛 {item['category']}" if item["category"] == "虫害" else f"🍃 {item['category']}")

                            st.markdown(f"**症状：** {item['symptoms']}")
                            st.markdown(f"**防治：** {item['treatment']}")

                # Advisory summary
                if results:
                    best = results[0][1]
                    st.subheader("💡 初步建议")
                    st.info(
                        f"系统判断该图片与 **{best['name']}**（{best['crop']}{best['category']}）最为相似，"
                        f"相似度 **{results[0][0]*100:.1f}%**。\n\n"
                        f"建议结合田间实际情况进一步确认，参考防治方案：**{best['treatment']}**"
                    )

            except Exception as e:
                st.error(f"搜索失败: {e}")

# ─── Footer ───────────────────────────────────────────────────────────────────
st.divider()
st.caption("病虫害以图搜图 · 基于 Qwen3-VL-Embedding · 结果仅供参考，请结合田间实际情况判断")