"""
Dify Knowledge Base Configuration Script
=========================================
Apply recommended settings to Dify KB.

Usage:
  python configure_dify_kb.py

Environment:
  DIFY_API_KEY: Dify dataset API Key
  DIFY_BASE_URL: Dify server URL (default http://dify.bianwoyou.cn)
"""
import os
import sys
import json
import requests

DIFY_BASE_URL = os.environ.get("DIFY_BASE_URL", "http://dify.bianwoyou.cn")
API_KEY = os.environ.get("DIFY_API_KEY", "")
DATASET_ID = "3ff939b3-8686-44f6-8ef5-65b1e53b55d3"

RECOMMENDED_CONFIG = {
    "retrieval_model": {
        "search_method": "hybrid_search",
        "reranking_enable": True,
        "reranking_mode": "weighted_score",
        "weights": {
            "vector_weight": 0.6,
            "keyword_weight": 0.4
        },
        "top_k": 6,
        "score_threshold_enabled": True,
        "score_threshold": 0.5
    }
}


def get_dataset_info():
    headers = {"Authorization": f"Bearer {API_KEY}"}
    url = f"{DIFY_BASE_URL}/v1/datasets/{DATASET_ID}"
    resp = requests.get(url, headers=headers, timeout=30)
    if resp.status_code == 200:
        return resp.json()
    print(f"  [FAIL] Get dataset info failed: {resp.status_code}")
    return None


def get_retrieval_config():
    headers = {"Authorization": f"Bearer {API_KEY}"}
    url = f"{DIFY_BASE_URL}/v1/datasets/{DATASET_ID}/retrieval-model"
    resp = requests.get(url, headers=headers, timeout=30)
    if resp.status_code == 200:
        return resp.json()
    print(f"  [FAIL] Get retrieval config failed: {resp.status_code}")
    return None


def update_retrieval_config():
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }
    url = f"{DIFY_BASE_URL}/v1/datasets/{DATASET_ID}/retrieval-model"
    resp = requests.post(url, headers=headers, json=RECOMMENDED_CONFIG, timeout=30)
    if resp.status_code in (200, 201):
        print("  [OK] Retrieval config updated successfully!")
        return True
    else:
        print(f"  [FAIL] Update failed: {resp.status_code} - {resp.text[:200]}")
        return False


def list_documents():
    headers = {"Authorization": f"Bearer {API_KEY}"}
    url = f"{DIFY_BASE_URL}/v1/datasets/{DATASET_ID}/documents?page=1&page_size=50"
    resp = requests.get(url, headers=headers, timeout=30)
    if resp.status_code == 200:
        docs = resp.json().get("data", [])
        print(f"\nDocuments ({len(docs)} total):")
        print(f"{'Name':<40} {'Status':<15} {'Words':<8}")
        print("-" * 65)
        for doc in docs:
            name = doc.get("name", "?")
            status = doc.get("indexing_status", "?")
            words = doc.get("word_count", 0)
            print(f"{name:<40} {status:<15} {words:<8}")
        return docs
    else:
        print(f"  [FAIL] List documents failed: {resp.status_code}")
        return []


def main():
    if not API_KEY:
        print("Error: Set DIFY_API_KEY environment variable")
        sys.exit(1)

    print("=" * 60)
    print("Dify Knowledge Base Configurator")
    print("=" * 60)
    print(f"Dataset ID: {DATASET_ID}")
    print(f"API URL:    {DIFY_BASE_URL}")
    print()

    # Step 1: Connectivity
    print("[1/4] Checking connectivity...")
    info = get_dataset_info()
    if info:
        name = info.get("name", "?")
        doc_count = info.get("document_count", "?")
        method = info.get("retrieval_model_dict", {}).get("search_method", "?")
        top_k = info.get("retrieval_model_dict", {}).get("top_k", "?")
        print(f"  [OK] KB: {name}")
        print(f"       Documents: {doc_count}")
        print(f"       Current: method={method}, top_k={top_k}")
    else:
        print("  [FAIL] Cannot reach Dify API")
        return

    # Step 2: Update config
    print()
    print("[2/4] Updating retrieval config...")
    print(f"       Target: method=hybrid_search, top_k=6, rerank=True")
    update_retrieval_config()

    # Step 3: Verify
    print()
    print("[3/4] Verifying config...")
    updated = get_retrieval_config()
    if updated:
        print(f"  [OK] Config applied")
        print(f"       {json.dumps(updated, ensure_ascii=False)[:200]}")
    else:
        print("  [WARN] Could not verify config (API may not support direct read)")

    # Step 4: List documents
    print()
    print("[4/4] Document list...")
    list_documents()

    print()
    print("=" * 60)
    print("Recommended settings (from design doc):")
    print("  Search: hybrid (vector 60% + keyword 40%)")
    print("  Top-K:  6")
    print("  Threshold: 0.5")
    print("  Rerank: enabled (weighted score)")
    print("=" * 60)


if __name__ == "__main__":
    main()