#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
将 vector_chunks.json 向量化并上传到 Milvus 和 PostgreSQL
使用中转站的 OpenAI 兼容 API
"""

import argparse
import json
import time
from pathlib import Path
from typing import List, Dict

import psycopg2
from psycopg2.extras import execute_values
from pymilvus import (
    connections,
    Collection,
    FieldSchema,
    CollectionSchema,
    DataType,
    utility,
)
from openai import OpenAI

# ===================== 配置 =====================
# 中转站配置
RELAY_BASE_URL = "http://6.86.80.4:30080/v1"
RELAY_API_KEY = "sk-5HeY7gfSIlyZMacfuXOf5cphpymsNqufEu1ou4U3avbULcyY"
EMBEDDING_MODEL = "text-embedding-v3"  # 中转站支持的 embedding 模型

# Milvus 配置
MILVUS_HOST = "localhost"
MILVUS_PORT = "19530"
COLLECTION_NAME = "regulation_chunks"

# PostgreSQL 配置
PG_HOST = "6.86.80.10"
PG_PORT = 5432
PG_USER = "postgresql"
PG_PASSWORD = "postgresql123456"
PG_DATABASE = "postgres"


# ===================== Embedding =====================
def get_openai_client(api_key: str, base_url: str) -> OpenAI:
    """创建 OpenAI 客户端连接到中转站"""
    return OpenAI(api_key=api_key, base_url=base_url)


def get_embeddings_batch(client: OpenAI, texts: List[str], batch_size: int = 10) -> List[List[float]]:
    """批量获取文本向量"""
    all_embeddings = []

    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        print(f"Embedding batch {i // batch_size + 1}/{(len(texts) - 1) // batch_size + 1}...")

        response = client.embeddings.create(
            model=EMBEDDING_MODEL,
            input=batch,
        )

        embeddings = [item.embedding for item in response.data]
        all_embeddings.extend(embeddings)

    return all_embeddings


# ===================== Milvus =====================
def init_milvus(host: str, port: str):
    connections.connect("default", host=host, port=port)
    print(f"已连接 Milvus: {host}:{port}")


def create_collection(name: str, dim: int) -> Collection:
    """创建或获取 collection"""
    if utility.has_collection(name):
        print(f"Collection '{name}' 已存在，删除重建")
        utility.drop_collection(name)

    fields = [
        FieldSchema(name="chunk_id", dtype=DataType.VARCHAR, max_length=64, is_primary=True),
        FieldSchema(name="doc_id", dtype=DataType.VARCHAR, max_length=128),
        FieldSchema(name="doc_title", dtype=DataType.VARCHAR, max_length=512),
        FieldSchema(name="chunk_index", dtype=DataType.INT64),
        FieldSchema(name="semantic_id", dtype=DataType.VARCHAR, max_length=64),
        FieldSchema(name="chunk_type", dtype=DataType.VARCHAR, max_length=32),
        FieldSchema(name="page_start", dtype=DataType.INT64),
        FieldSchema(name="page_end", dtype=DataType.INT64),
        FieldSchema(name="section_title", dtype=DataType.VARCHAR, max_length=512),
        FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=2048),
        FieldSchema(name="source_ids", dtype=DataType.VARCHAR, max_length=4096),  # JSON 字符串
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=dim),
    ]

    schema = CollectionSchema(fields, description="法规文档检索 chunks")
    collection = Collection(name, schema)

    # 创建向量索引（IVF_FLAT，适合中小规模）
    index_params = {
        "metric_type": "COSINE",
        "index_type": "IVF_FLAT",
        "params": {"nlist": 128},
    }
    collection.create_index("embedding", index_params)
    print(f"Collection '{name}' 创建完成，索引已建立")

    return collection


def insert_chunks(collection: Collection, chunks: List[Dict], embeddings: List[List[float]]):
    """插入 chunks 到 Milvus"""
    data = [
        [c["chunk_id"] for c in chunks],
        [c["doc_id"] for c in chunks],
        [c["doc_title"] for c in chunks],
        [c["chunk_index"] for c in chunks],
        [c["semantic_id"] for c in chunks],
        [c["chunk_type"] for c in chunks],
        [c["page_start"] for c in chunks],
        [c["page_end"] for c in chunks],
        [c["section_title"] for c in chunks],
        [c["text"] for c in chunks],
        [json.dumps(c.get("source_ids", [])) for c in chunks],  # JSON 字符串
        embeddings,
    ]

    collection.insert(data)
    collection.flush()
    print(f"已插入 {len(chunks)} 个 chunks")


def load_collection(collection: Collection):
    """加载 collection 到内存（搜索前必须）"""
    collection.load()
    print(f"Collection 已加载到内存")


# ===================== PostgreSQL =====================
def get_pg_connection(host: str, port: int, user: str, password: str, database: str):
    """获取 PostgreSQL 连接"""
    conn = psycopg2.connect(
        host=host,
        port=port,
        user=user,
        password=password,
        database=database,
    )
    print(f"已连接 PostgreSQL: {host}:{port}/{database}")
    return conn


def insert_chunks_to_pg(conn, chunks: List[Dict], doc_data: Dict):
    """插入 chunks 和相关数据到 PostgreSQL"""
    cursor = conn.cursor()

    try:
        # 1. 插入文档
        cursor.execute("""
            INSERT INTO documents (doc_id, title, standard_number, upload_time)
            VALUES (%s, %s, %s, NOW())
            ON CONFLICT (doc_id) DO UPDATE SET title = EXCLUDED.title, updated_at = NOW()
        """, (doc_data["doc_id"], doc_data["doc_title"], doc_data.get("standard_number")))

        # 2. 插入语义块
        semantic_blocks = doc_data.get("semantic_blocks", [])
        if semantic_blocks:
            block_rows = [
                (
                    doc_data["doc_id"],
                    block["semantic_id"],
                    block["block_type"],
                    block["page_start"],
                    block["page_end"],
                    block.get("section_title"),
                    block.get("section_level"),
                    json.dumps(block.get("source_ids", [])),
                    block["text"],
                )
                for block in semantic_blocks
            ]
            execute_values(
                cursor,
                """
                INSERT INTO semantic_blocks
                (doc_id, semantic_id, block_type, page_start, page_end, section_title, section_level, source_ids, text)
                VALUES %s
                ON CONFLICT (doc_id, semantic_id) DO UPDATE SET text = EXCLUDED.text
                """,
                block_rows,
            )
            print(f"已插入 {len(semantic_blocks)} 个语义块")

        # 3. 插入向量块元数据
        chunk_rows = [
            (
                doc_data["doc_id"],
                chunk["chunk_id"],
                chunk["semantic_id"],
                chunk["chunk_index"],
                chunk.get("piece_index"),
                chunk["page_start"],
                chunk["page_end"],
                chunk.get("section_title"),
                chunk["text"],
                json.dumps(chunk.get("source_ids", [])),
            )
            for chunk in chunks
        ]
        execute_values(
            cursor,
            """
            INSERT INTO vector_chunks
            (doc_id, chunk_id, semantic_id, chunk_index, piece_index, page_start, page_end, section_title, text, source_ids)
            VALUES %s
            ON CONFLICT (doc_id, chunk_id) DO UPDATE SET text = EXCLUDED.text
            """,
            chunk_rows,
        )
        print(f"已插入 {len(chunks)} 个向量块元数据")

        conn.commit()
        print("PostgreSQL 数据插入完成")

    except Exception as e:
        conn.rollback()
        raise e
    finally:
        cursor.close()


# ===================== 主流程 =====================
def load_data(file_path: Path) -> Dict:
    """加载 vector_chunks.json，返回完整数据"""
    data = json.loads(file_path.read_text(encoding="utf-8"))
    return data


def upload_to_milvus_and_pg(
    chunks_file: str,
    api_key: str,
    base_url: str,
    milvus_host: str,
    milvus_port: str,
    collection_name: str,
    batch_size: int,
    pg_host: str,
    pg_port: int,
    pg_user: str,
    pg_password: str,
    pg_database: str,
):
    # 1. 加载完整数据
    chunks_path = Path(chunks_file).expanduser().resolve()
    if not chunks_path.exists():
        raise FileNotFoundError(f"文件不存在: {chunks_path}")

    data = load_data(chunks_path)
    chunks = data.get("vector_chunks", [])
    if not chunks:
        raise ValueError("vector_chunks 为空")
    print(f"加载 {len(chunks)} 个 chunks")

    # 2. 初始化连接
    client = get_openai_client(api_key, base_url)
    init_milvus(milvus_host, milvus_port)
    pg_conn = get_pg_connection(pg_host, pg_port, pg_user, pg_password, pg_database)

    # 3. 获取 embeddings
    texts = [c["embedding_text"] for c in chunks]
    embeddings = get_embeddings_batch(client, texts, batch_size)
    print(f"生成 {len(embeddings)} 个向量")

    # 4. 获取 embedding 维度
    embedding_dim = len(embeddings[0])
    print(f"Embedding 维度: {embedding_dim}")

    # 5. 创建 collection 并插入 Milvus
    collection = create_collection(collection_name, embedding_dim)
    insert_chunks(collection, chunks, embeddings)
    load_collection(collection)

    # 6. 插入 PostgreSQL
    insert_chunks_to_pg(pg_conn, chunks, data)

    # 7. 关闭连接
    pg_conn.close()

    print("上传完成！")


# ===================== CLI =====================
def main():
    parser = argparse.ArgumentParser(description="将 vector_chunks 向量化并上传到 Milvus 和 PostgreSQL")
    parser.add_argument("chunks_file", help="vector_chunks.json 文件路径")
    parser.add_argument("--api-key", default=RELAY_API_KEY, help="中转站 API Key")
    parser.add_argument("--base-url", default=RELAY_BASE_URL, help="中转站 Base URL")
    parser.add_argument("--milvus-host", default=MILVUS_HOST, help="Milvus host")
    parser.add_argument("--milvus-port", default=MILVUS_PORT, help="Milvus port")
    parser.add_argument("--collection", default=COLLECTION_NAME, help="Milvus collection 名称")
    parser.add_argument("--batch-size", type=int, default=10, help="Embedding 批量大小（中转站限制最大10）")
    parser.add_argument("--pg-host", default=PG_HOST, help="PostgreSQL host")
    parser.add_argument("--pg-port", type=int, default=PG_PORT, help="PostgreSQL port")
    parser.add_argument("--pg-user", default=PG_USER, help="PostgreSQL user")
    parser.add_argument("--pg-password", default=PG_PASSWORD, help="PostgreSQL password")
    parser.add_argument("--pg-database", default=PG_DATABASE, help="PostgreSQL database")
    args = parser.parse_args()

    upload_to_milvus_and_pg(
        chunks_file=args.chunks_file,
        api_key=args.api_key,
        base_url=args.base_url,
        milvus_host=args.milvus_host,
        milvus_port=args.milvus_port,
        collection_name=args.collection,
        batch_size=args.batch_size,
        pg_host=args.pg_host,
        pg_port=args.pg_port,
        pg_user=args.pg_user,
        pg_password=args.pg_password,
        pg_database=args.pg_database,
    )


if __name__ == "__main__":
    main()