#!/usr/bin/env python3
"""
Fix DeepSeek API configuration to ensure it works without fallback models
"""

import os
import sys
import subprocess
import requests
import json
import time
from pathlib import Path

# Add current directory to Python path
current_dir = Path(__file__).parent
sys.path.insert(0, str(current_dir))

def verify_deepseek_direct():
    """Test DeepSeek API directly"""
    print("🔍 Testing DeepSeek API directly...")
    
    try:
        import openai
        from openai import OpenAI
        
        # Test with direct API call
        client = OpenAI(
            api_key="sk-55f6e57f1d834b0e93ceaf98cc2cb715",
            base_url="https://api.deepseek.com/v1"
        )
        
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=[{"role": "user", "content": "Test message to verify API connectivity"}],
            max_tokens=50
        )
        
        print(f"✅ DeepSeek API direct test successful")
        print(f"🤖 Response: {response.choices[0].message.content}")
        return True
        
    except Exception as e:
        print(f"❌ DeepSeek API direct test failed: {e}")
        return False

def check_server_config():
    """Check current server configuration"""
    print("\n📋 Checking current server configuration...")
    
    # Check if server is running
    try:
        response = requests.get("http://localhost:3015/api/health", timeout=5)
        if response.status_code == 200:
            print("✅ Server is running")
            return True
        else:
            print(f"❌ Server returned status: {response.status_code}")
            return False
    except:
        print("❌ Server is not running")
        return False

def create_fixed_server_script():
    """Create a fixed server startup script that forces DeepSeek configuration"""
    script_content = '''#!/usr/bin/env python3
"""
Fixed LightRAG server startup with forced DeepSeek API configuration
"""

import os
import sys
from pathlib import Path

# Force set environment variables for DeepSeek
os.environ["LLM_BINDING"] = "openai"
os.environ["LLM_MODEL"] = "deepseek-chat"
os.environ["LLM_BINDING_HOST"] = "https://api.deepseek.com/v1"
os.environ["LLM_BINDING_API_KEY"] = "sk-55f6e57f1d834b0e93ceaf98cc2cb715"
os.environ["EMBEDDING_BINDING"] = "ollama"
os.environ["EMBEDDING_MODEL"] = "snowflake-arctic-embed2:latest"
os.environ["EMBEDDING_DIM"] = "1024"

# Force storage backends
os.environ["LIGHTRAG_KV_STORAGE"] = "RedisKVStorage"
os.environ["LIGHTRAG_VECTOR_STORAGE"] = "QdrantVectorDBStorage"
os.environ["LIGHTRAG_DOC_STATUS_STORAGE"] = "PGDocStatusStorage"
os.environ["LIGHTRAG_GRAPH_STORAGE"] = "Neo4JStorage"

# Database connections
os.environ["REDIS_URI"] = "redis://localhost:6379"
os.environ["NEO4J_URI"] = "bolt://localhost:7687"
os.environ["NEO4J_USERNAME"] = "neo4j"
os.environ["NEO4J_PASSWORD"] = "jleu1212"
os.environ["QDRANT_URL"] = "http://localhost:6333"
os.environ["POSTGRES_HOST"] = "localhost"
os.environ["POSTGRES_PORT"] = "5432"
os.environ["POSTGRES_USER"] = "jleu3482"
os.environ["POSTGRES_PASSWORD"] = "jleu1212"
os.environ["POSTGRES_DATABASE"] = "rag_anything"

print("🔧 Forced DeepSeek API configuration:")
print(f"   LLM: {os.environ['LLM_BINDING']} - {os.environ['LLM_MODEL']}")
print(f"   Embedding: {os.environ['EMBEDDING_BINDING']} - {os.environ['EMBEDDING_MODEL']}")
print(f"   Storage: {os.environ['LIGHTRAG_KV_STORAGE']}, {os.environ['LIGHTRAG_VECTOR_STORAGE']}")

# Now start the server
lightrag_dir = Path(__file__).parent / "LightRAG-main"
os.chdir(lightrag_dir)

import subprocess
cmd = [
    sys.executable, "-m", "lightrag.api.lightrag_server",
    "--port", "3015",
    "--host", "0.0.0.0",
    "--working-dir", "rag_storage",
    "--input-dir", "../inputs",
    "--key", "jleu1212",
    "--auto-scan-at-startup",
    "--llm-binding", os.environ["LLM_BINDING"],
    "--embedding-binding", os.environ["EMBEDDING_BINDING"],
    "--rerank-binding", "jina"
]

print(f"🚀 Starting server with command: {' '.join(cmd)}")
subprocess.run(cmd, check=True)
'''

    script_path = current_dir / "start_deepseek_forced.py"
    with open(script_path, 'w') as f:
        f.write(script_content)
    
    print(f"✅ Created forced configuration script: {script_path}")
    return script_path

def test_retrieval_with_deepseek():
    """Test retrieval with DeepSeek API"""
    print("\n🔍 Testing retrieval with DeepSeek API...")
    
    # Wait for server to be ready
    print("⏳ Waiting for server to be ready...")
    time.sleep(10)
    
    try:
        headers = {
            'X-API-Key': 'jleu1212',
            'Content-Type': 'application/json'
        }
        data = {
            'query': 'railway transportation system',
            'top_k': 3,
            'use_llm': True
        }
        
        response = requests.post(
            "http://localhost:3015/api/search",
            headers=headers,
            json=data,
            timeout=30
        )
        
        if response.status_code == 200:
            results = response.json()
            print("✅ Search request successful")
            
            if results.get('llm_response'):
                print(f"🤖 DeepSeek LLM Response: {results['llm_response'][:200]}...")
                print("✅ DeepSeek API is working for retrieval!")
                return True
            else:
                print("⚠️ No LLM response in results")
                if results.get('results'):
                    print(f"📊 Vector search results: {len(results['results'])} found")
                    return True
                else:
                    print("❌ No search results at all")
                    return False
        else:
            print(f"❌ Search failed: {response.status_code}")
            print(f"Error: {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ Search test error: {e}")
        return False

def main():
    """Main function to fix DeepSeek API configuration"""
    print("=" * 60)
    print("🔧 FIXING DEEPSEEK API CONFIGURATION")
    print("=" * 60)
    
    # Step 1: Test DeepSeek API directly
    if not verify_deepseek_direct():
        print("❌ DeepSeek API is not accessible")
        return False
    
    # Step 2: Create fixed server script
    script_path = create_fixed_server_script()
    
    print(f"\n🚀 To start server with fixed DeepSeek configuration:")
    print(f"   python {script_path.name}")
    
    print("\n📋 Configuration Summary:")
    print("   ✅ DeepSeek API: openai/deepseek-chat")
    print("   ✅ Embedding: ollama/snowflake-arctic-embed2:latest")
    print("   ✅ Storage: Redis + Neo4j + Qdrant + PostgreSQL")
    print("   ✅ No fallback models enabled")
    
    print("\n🎯 Next steps:")
    print("1. Run: python start_deepseek_forced.py")
    print("2. Wait for server to start")
    print("3. Test with: python test_deepseek_retrieval.py")
    print("4. Verify DeepSeek responses in search results")
    
    return True

if __name__ == "__main__":
    success = main()
    exit(0 if success else 1)