#!/usr/bin/env python3
"""
Start LightRAG server with fixed DeepSeek API configuration
Ensures DeepSeek API is used without fallback models
"""

import os
import sys
import subprocess
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 load_env_vars():
    """Load environment variables from .env file"""
    env_file = current_dir / ".env"
    if env_file.exists():
        print(f"📁 Loading environment from {env_file}")
        with open(env_file, 'r') as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    os.environ[key.strip()] = value.strip()
                    print(f"   ✅ Set {key.strip()}")
    
    # Ensure critical DeepSeek variables are set
    required_vars = {
        'LLM_BINDING': 'openai',
        'LLM_MODEL': 'deepseek-chat',
        'LLM_BINDING_HOST': 'https://api.deepseek.com/v1',
        'LLM_BINDING_API_KEY': 'sk-55f6e57f1d834b0e93ceaf98cc2cb715'
    }
    
    for var, default in required_vars.items():
        if var not in os.environ:
            os.environ[var] = default
            print(f"   ⚠️  Set default {var}={default}")

def verify_deepseek_config():
    """Verify DeepSeek API configuration"""
    print("\n🔍 Verifying DeepSeek API configuration...")
    
    required_config = {
        'LLM_BINDING': 'openai',
        'LLM_MODEL': 'deepseek-chat', 
        'LLM_BINDING_HOST': 'https://api.deepseek.com/v1',
        'EMBEDDING_BINDING': 'ollama',
        'EMBEDDING_MODEL': 'snowflake-arctic-embed2:latest',
        'EMBEDDING_DIM': '1024'
    }
    
    all_good = True
    for key, expected in required_config.items():
        actual = os.environ.get(key, 'NOT_SET')
        if actual == expected:
            print(f"   ✅ {key}: {actual}")
        else:
            print(f"   ❌ {key}: {actual} (expected: {expected})")
            all_good = False
    
    return all_good

def start_server():
    """Start the LightRAG server with proper configuration"""
    print("\n🚀 Starting LightRAG server with DeepSeek API...")
    
    # Change to LightRAG directory
    lightrag_dir = current_dir / "LightRAG-main"
    os.chdir(lightrag_dir)
    
    # Build command with explicit environment variables
    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.get('LLM_BINDING', 'openai'),
        "--embedding-binding", os.environ.get('EMBEDDING_BINDING', 'ollama'),
        "--rerank-binding", "jina"
    ]
    
    print(f"   Command: {' '.join(cmd)}")
    print(f"   DeepSeek API Key: {os.environ.get('LLM_BINDING_API_KEY', 'NOT_SET')[:10]}...")
    
    # Start the server
    try:
        subprocess.run(cmd, check=True)
    except subprocess.CalledProcessError as e:
        print(f"❌ Server failed to start: {e}")
        return False
    except KeyboardInterrupt:
        print("\n🛑 Server stopped by user")
        return True
    
    return True

def test_deepseek_connection():
    """Test DeepSeek API connection"""
    print("\n🔌 Testing DeepSeek API connection...")
    
    try:
        import openai
        from openai import OpenAI
        
        # Configure OpenAI client for DeepSeek
        client = OpenAI(
            api_key=os.environ.get('LLM_BINDING_API_KEY'),
            base_url=os.environ.get('LLM_BINDING_HOST')
        )
        
        # Test a simple completion
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=[{"role": "user", "content": "Hello, are you working?"}],
            max_tokens=50
        )
        
        print(f"   ✅ DeepSeek API test successful")
        print(f"   🤖 Response: {response.choices[0].message.content}")
        return True
        
    except Exception as e:
        print(f"   ❌ DeepSeek API test failed: {e}")
        return False

def main():
    """Main function to start server with fixed DeepSeek configuration"""
    print("=" * 60)
    print("🔧 DEEPSEEK API FIXED CONFIGURATION")
    print("=" * 60)
    
    # Step 1: Load environment variables
    load_env_vars()
    
    # Step 2: Verify configuration
    if not verify_deepseek_config():
        print("❌ Configuration verification failed")
        return False
    
    # Step 3: Test DeepSeek connection
    if not test_deepseek_connection():
        print("❌ DeepSeek API connection test failed")
        return False
    
    # Step 4: Start server
    print("\n" + "=" * 40)
    print("🚀 STARTING SERVER")
    print("=" * 40)
    
    return start_server()

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