#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Быстрая проверка модели OpenCode на способность извлекать сущности (как требует LightRAG).

Использование:
    python test_model.py <model_name>

Примеры:
    python test_model.py deepseek-v4-flash-free
    python test_model.py nemotron-3-super-free
    python test_model.py minimax-m2.5-free
"""

import os
import sys
import asyncio
from dotenv import load_dotenv
from openai import AsyncOpenAI

load_dotenv()

TEST_PROMPT = """
Извлеки сущности и связи из текста. Ответь строго в формате:

### Сущности:
1. entity_name (тип)

### Связи:
1. entity1 -> relationship -> entity2

Текст:
Страница 5 чертежа. 3-й этаж здания. Оси А и Б. Размер 5400 мм между осями. Квартира 101 на этаже 3.
"""

async def test_model(model: str):
    base_url = os.environ.get("OPENCODE_URL", "https://opencode.ai/zen/v1")
    api_key = os.environ.get("OPENCODE_API_KEY", "")
    
    if not api_key:
        print("[ERR] OPENCODE_API_KEY не найден в .env")
        sys.exit(1)
    
    print(f"[TEST] Модель: {model}")
    print(f"[TEST] URL: {base_url}")
    print(f"[TEST] Отправка запроса...\n")
    
    client = AsyncOpenAI(base_url=base_url, api_key=api_key)
    
    try:
        response = await client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "Ты — анализатор чертежей. Извлекай структурированные сущности и связи."},
                {"role": "user", "content": TEST_PROMPT}
            ],
            temperature=0.3,
            max_tokens=1024,
        )
        
        result = response.choices[0].message.content
        print("=" * 60)
        print("ОТВЕТ МОДЕЛИ:")
        print("=" * 60)
        print(result)
        print("=" * 60)
        
        # Простая проверка
        has_entities = "### Сущности" in result or "entity" in result.lower() or "сущности" in result.lower()
        has_relations = "### Связи" in result or "relationship" in result.lower() or "связи" in result.lower()
        
        print(f"\n[РЕЗУЛЬТАТ]")
        print(f"  Найдены сущности: {'ДА' if has_entities else 'НЕТ'}")
        print(f"  Найдены связи: {'ДА' if has_relations else 'НЕТ'}")
        
        if has_entities and has_relations:
            print(f"\n✅ Модель {model} ПОДХОДИТ для LightRAG")
        else:
            print(f"\n❌ Модель {model} НЕ подходит для LightRAG (не следует формату)")
            print("   Попробуйте более мощную модель (GPT-4o, Claude, Llama 3.1 70B)")
            
    except Exception as e:
        print(f"\n[ERR] Ошибка API: {e}")
        print("   Возможно, модель недоступна или требует другого имени")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python test_model.py <model_name>")
        print("Examples:")
        print("  python test_model.py deepseek-v4-flash-free")
        print("  python test_model.py nemotron-3-super-free")
        print("  python test_model.py minimax-m2.5-free")
        sys.exit(1)
    
    asyncio.run(test_model(sys.argv[1]))