"""Basic usage examples for DataBuilder Logic SDK."""

import os
from collections.abc import Callable

from dotenv import load_dotenv

from databuilder_logic_sdk.core import config, create_embedding_service, create_ontology_service, create_rerank_service
from databuilder_logic_sdk.utils import logger


def _run_section(title: str, handler: Callable[[], None]) -> None:
    logger.info(f"\n=== {title} ===")
    try:
        handler()
    except Exception as exc:
        logger.error(f"{title} failed: {exc}")


def main() -> None:
    """Demonstrate basic usage of the SDK."""

    load_dotenv()

    logger.debug(f"env {os.getenv('ServiceConfig')}")
    logger.debug(f"parsed config: {config}")

    def run_embedding() -> None:
        embedding_service = create_embedding_service("text-embed_7b_bf16")
        texts = ["这是一个测试文本。", "另一个测试"]
        response = embedding_service.embed(texts)
        for result in response.results:
            logger.info(f"Index {result.index}: embedding length = {len(result.embedding)}")

        bge_zh = create_embedding_service(
            "bge-large-zh",
            endpoints=["http://qianfan.baidubce.com"],
            api_key=os.getenv("BCE_API_KEY"),
        )
        bge_response = bge_zh.embed(["测试BGE中文模型", "再测试BGE"])
        for result in bge_response.results:
            logger.info(f"BGE-ZH Index {result.index}: embedding length = {len(result.embedding)}")

        bge_en = create_embedding_service(
            "bge-large-en",
            endpoints=["http://qianfan.baidubce.com"],
            api_key=os.getenv("BCE_API_KEY"),
        )
        bge_en_response = bge_en.embed(["Testing BGE English model", "Testing BGE again"])
        for result in bge_en_response.results:
            logger.info(f"BGE-EN Index {result.index}: embedding length = {len(result.embedding)}")

    def run_rerank() -> None:
        rerank_service = create_rerank_service("bce-reranker-base")
        query = "北京人工智能"
        documents = [
            "致力于创建能够执行通常需要人类智能的任务的系统。",
            "机器学习、自然语言处理、计算机视觉等领域。",
            "北京天气。",
        ]
        rerank_response = rerank_service.rerank(
            query=query,
            documents=documents,
            top_k=2,
        )
        logger.info(f"Rerank query: {query}")
        logger.info("Rerank Results:")
        # Since rerank_response is a pydantic model, you can convert it to a python dict if needed
        # rerank_response = rerank_response.model_dump()
        for rerank_result in rerank_response.results:
            logger.info(f"Document: {rerank_result.document}")
            logger.info(f"Score: {rerank_result.relevance_score}")
            logger.info(f"Index: {rerank_result.index}")
            logger.info("---")

        logger.info("Sorted Documents:")
        sorted_results = rerank_service.sort(documents, rerank_response.results)
        for result in sorted_results:
            logger.info(f"  `{result}`")

    def run_ontology() -> None:
        service = create_ontology_service()
        dsl = {
            "ontology": "whb_demo01",
            "apiName": "datasearch0122",
            "limit": 5,
        }
        ontology_response = service.search(
            dsl=dsl,
        )
        logger.info(f"Ontology Response Code: {ontology_response.code}")
        if ontology_response.result and "data" in ontology_response.result:
            logger.info(f"Ontology Data Count: {len(ontology_response.result['data'])}")

    _run_section("Embedding Service Example", run_embedding)
    _run_section("Rerank Service Example", run_rerank)
    _run_section("Ontology Service Example", run_ontology)


if __name__ == "__main__":
    main()
