"""AIGC 短剧「生成 → 检索」一体化流水线（阿里云 Milvus AI Function）
步骤1 关键帧生图 → 步骤2 图生视频 → 步骤3 素材理解 → 步骤4 向量入库 → 步骤5 内容搜索
用到四个 AI Function：AI_IMAGE_EDIT / AI_VIDEO_EDIT / AI_MULTI_MODAL_GENERATE / AI_EMBEDDING
"""
from __future__ import annotations

import json
import time
from typing import Any
from urllib.error import HTTPError
from urllib.request import Request, urlopen

from pymilvus import DataType, Function, FunctionType, MilvusClient

# ==== 全局配置（REST 与 gRPC 同在 19530 端口）====
MILVUS_ENDPOINT = "c-xxx.milvus.aliyuncs.com:19530"
MILVUS_TOKEN = "root:xxx"
MILVUS_REST_BASE_URL = f"http://{MILVUS_ENDPOINT}"
MILVUS_URI = MILVUS_REST_BASE_URL


def post_json(path: str, body: dict[str, Any], timeout: int = 200) -> tuple[int, dict[str, Any]]:
    """步骤 1~3 共用的 REST 调用工具。"""
    request = Request(
        f"{MILVUS_REST_BASE_URL.rstrip('/')}{path}",
        data=json.dumps(body, ensure_ascii=False).encode("utf-8"),
        headers={"Authorization": f"Bearer {MILVUS_TOKEN}", "Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urlopen(request, timeout=timeout) as response:
            return response.status, json.loads(response.read().decode("utf-8"))
    except HTTPError as exc:
        return exc.code, json.loads(exc.read().decode("utf-8"))


# ========== 步骤 1：关键帧生图 AI_IMAGE_EDIT (wan2.7-image)　【对应正文·步骤1｜代码第 37–55 行】
IMAGE_URL = "https://dashscope.oss-cn-beijing.aliyuncs.com/images/dog_and_girl.jpeg"
PROMPT = (
    "雨夜，一位女子撑着伞站在便利店门口，暖色的店内灯光打在她脸上，"
    "电影级构图，保留画面主体。"
)
status, data = post_json(
    "/v2/vectordb/ai/image_edit",
    {
        "model_name": "wan2.7-image",
        "texts": [IMAGE_URL],
        "params": {"prompt": PROMPT, "n": 1, "size": "1024*1024", "watermark": False, "timeout_sec": 180},
    },
)
assert status == 200 and data.get("code") == 0, data
keyframe_url = data["data"]["output"]["outputs"][0]
print("\n" + "=" * 64)
print("步骤 1 | AI_IMAGE_EDIT 关键帧生图（wan2.7-image）")
print("=" * 64)
print(f"关键帧图片 URL：{keyframe_url}")


# ========== 步骤 2：图生视频 AI_VIDEO_EDIT (happyhorse-1.1-i2v)，异步任务　【对应正文·步骤2｜代码第 56–91 行】
VIDEO_MODEL = "happyhorse-1.1-i2v"  # 创建与查询必须使用同一模型名
# 用步骤1的关键帧作首帧（media.type=first_frame）
status, data = post_json(
    "/v2/vectordb/ai/video_edit",
    {
        "model_name": VIDEO_MODEL,
        "prompt": "镜头缓缓推进，雨丝飘落，暖光在伞面上轻轻闪动，"
                  "保留首帧图中的主体。",
        "media": [{"type": "first_frame", "url": keyframe_url}],
        "params": {"resolution": "720P", "audio_setting": "none", "watermark": False, "timeout_sec": 180},
    },
)
assert status == 200 and data.get("code") == 0, data
task_id = data["data"]["output"]["task_id"]
print("\n" + "=" * 64)
print("步骤 2 | AI_VIDEO_EDIT 图生视频（happyhorse-1.1-i2v，异步任务）")
print("=" * 64)
print(f"任务已创建  task_id = {task_id}")

# 轮询查询：tasks/describe 一次只查一个 task_id，带 provider 与相同 model_name
video_url = None
for attempt in range(60):  # 本地轮询上限，达到上限只记录 task_id，不判定服务端失败
    status, data = post_json(
        "/v2/vectordb/ai/tasks/describe",
        {"provider": "aliyun_milvus", "model_name": VIDEO_MODEL, "task_id": task_id},
    )
    task_status = data.get("data", {}).get("output", {}).get("task_status")
    print(f"  轮询 #{attempt:<2} ->  {task_status}")
    if task_status == "SUCCEEDED":
        video_url = data["data"]["output"]["video_url"]
        break
    if task_status in ("FAILED", "CANCELED"):
        raise RuntimeError(f"task ended in {task_status}: {json.dumps(data, ensure_ascii=False)}")
    time.sleep(10)
assert video_url, f"轮询达到本地上限，请稍后凭 task_id={task_id} 继续查询"
print(f"运动镜头 video_url：{video_url}")


# ========== 步骤 3：素材理解 AI_MULTI_MODAL_GENERATE (qwen3.7-plus)　【对应正文·步骤3｜代码第 92–117 行】
def describe_media(url: str, media_type: str, prompt: str) -> str:
    status, data = post_json(
        "/v2/vectordb/ai/multi_modal_generate",
        {
            "model_name": "qwen3.7-plus",
            "texts": [url],
            "params": {"media_type": media_type, "prompt": prompt, "temperature": 0},
        },
    )
    assert status == 200 and data.get("code") == 0, data
    return data["data"]["output"]["outputs"][0]


image_caption = describe_media(
    keyframe_url, "image",
    "用一句中文客观描述这张关键帧图片的主体、场景和色调，便于以文搜图。",
)
video_caption = describe_media(
    video_url, "video",
    "用一句中文客观描述这段视频的主体、动作和画面氛围，便于按镜头检索。",
)
print("\n" + "=" * 64)
print("步骤 3 | AI_MULTI_MODAL_GENERATE 素材理解（qwen3.7-plus）")
print("=" * 64)
print(f"[图片描述] {image_caption}")
print(f"[视频描述] {video_caption}")


# ========== 步骤 4：建多模态 Collection 并向量入库 AI_EMBEDDING (qwen3-vl-embedding)　【对应正文·步骤4｜代码第 118–168 行】
EMBED_MODEL = "qwen3-vl-embedding"
VECTOR_DIM = 2560  # qwen3-vl-embedding 多模态维度，向量字段维度必须与 dim 一致

client = MilvusClient(uri=MILVUS_URI, token=MILVUS_TOKEN)

collection_name = "aigc_assets_multimodal"
if client.has_collection(collection_name):
    client.drop_collection(collection_name)

schema = MilvusClient.create_schema(auto_id=True, enable_dynamic_field=False)
schema.add_field("id", DataType.INT64, is_primary=True)
schema.add_field("media_type", DataType.VARCHAR, max_length=16)     # image / video
schema.add_field("media_ref", DataType.VARCHAR, max_length=4096)    # 用于向量化的图片/视频 URL
schema.add_field("caption", DataType.VARCHAR, max_length=2048)      # 步骤3 生成的描述
schema.add_field("tags", DataType.VARCHAR, max_length=512)
schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=VECTOR_DIM)

# 写入 media_ref 时自动调用 qwen3-vl-embedding 生成 2560 维向量
schema.add_function(
    Function(
        name="embed_media",
        function_type=FunctionType.TEXTEMBEDDING,
        input_field_names=["media_ref"],
        output_field_names=["embedding"],
        params={
            "provider": "aliyun_milvus",
            "model_name": EMBED_MODEL,
            "dim": VECTOR_DIM,
            "is_multimodal": "true",
        },
    )
)

index_params = client.prepare_index_params()
index_params.add_index(field_name="embedding", index_type="AUTOINDEX", metric_type="COSINE")
client.create_collection(collection_name=collection_name, schema=schema, index_params=index_params)

# 把步骤1~3 的产出（url + caption + tags）向量化入库
client.insert(
    collection_name,
    [
        {"media_type": "image", "media_ref": keyframe_url, "caption": image_caption, "tags": "关键帧,生成图"},
        {"media_type": "video", "media_ref": video_url, "caption": video_caption, "tags": "运动镜头,生成视频"},
    ],
)
client.flush(collection_name)
client.load_collection(collection_name)
print("\n" + "=" * 64)
print("步骤 4 | AI_EMBEDDING 多模态向量入库（qwen3-vl-embedding, 2560 维）")
print("=" * 64)
print(f"图 + 视频已向量化入库，Collection：{collection_name}")


# ========== 步骤 5：内容搜索 以文搜图 / 以文搜视频 / 以图搜视频　【对应正文·步骤5｜代码第 169–196 行】
def search_assets(query: str, top_k: int = 5, media_type: str | None = None):
    """query 可以是一段中文/英文文案，也可以是一张查询图片的 URL。"""
    filter_expr = f'media_type == "{media_type}"' if media_type else ""
    results = client.search(
        collection_name=collection_name,
        data=[query],                 # 文本或图片 URL，均由 qwen3-vl-embedding 向量化
        anns_field="embedding",
        limit=top_k,
        filter=filter_expr,
        output_fields=["media_type", "media_ref", "caption", "tags"],
    )
    for rank, hit in enumerate(results[0], 1):
        e = hit["entity"]
        url = e["media_ref"].split("?")[0]   # 去掉 OSS 签名参数，展示更干净
        print(f"  {rank}. [相似度 {hit['distance']:.4f}] [{e['media_type']}] {e['caption'][:42]}")
        print(f"       {url}")


print("\n" + "=" * 64)
print("步骤 5 | 多模态内容搜索（qwen3-vl-embedding 统一向量空间）")
print("=" * 64)
# 以文搜图：限定只搜图片
print("[5.1] 以文搜图：")
search_assets("室内暖光下的条纹毛衣造型", media_type="image")

# 以文搜视频：限定只搜视频
print("\n[5.2] 以文搜视频：")
search_assets("角色近景，镜头缓慢推进的概念片", media_type="video")

# 以图搜视频：用一张关键帧图片 URL 去召回风格相近的视频素材
print("\n[5.3] 以图搜视频（用关键帧图召回同源视频）：")
search_assets(keyframe_url, media_type="video")