#!/usr/bin/env python3
"""Diagnostic script for HF Space deployment issues"""

import os
import sys
from pathlib import Path

def check_files():
    """Check if all required files exist"""
    print("=== File Check ===")
    
    required_files = [
        "config/settings.py",
        "src/retrieval/semantic_retriever.py", 
        "src/generation/response_generator.py",
        "data/vector_store/embeddings.npy",
        "data/vector_store/faiss_index.bin",
        "data/vector_store/chunks.pkl",
        "data/vector_store/bm25.pkl",
        "data/vector_store/relationships.pkl"
    ]
    
    for file_path in required_files:
        if os.path.exists(file_path):
            print(f"✅ {file_path}")
        else:
            print(f"❌ {file_path} - MISSING")

def check_imports():
    """Check if imports work"""
    print("\n=== Import Check ===")
    
    try:
        from config.settings import config
        print("✅ Config imported")
        print(f"   - Embedding model: {config.EMBEDDING_MODEL}")
        print(f"   - Vector store path: {config.VECTOR_STORE_PATH}")
    except Exception as e:
        print(f"❌ Config import failed: {e}")
        return
    
    try:
        from src.retrieval.semantic_retriever import SemanticRetriever
        print("✅ SemanticRetriever imported")
    except Exception as e:
        print(f"❌ SemanticRetriever import failed: {e}")
    
    try:
        from src.generation.response_generator import ResponseGenerator
        print("✅ ResponseGenerator imported")
    except Exception as e:
        print(f"❌ ResponseGenerator import failed: {e}")
    
    try:
        from sentence_transformers import SentenceTransformer
        print("✅ SentenceTransformers imported")
    except Exception as e:
        print(f"❌ SentenceTransformers import failed: {e}")
    
    try:
        import faiss
        print("✅ FAISS imported")
    except Exception as e:
        print(f"❌ FAISS import failed: {e}")

def check_vector_store():
    """Check vector store files"""
    print("\n=== Vector Store Check ===")
    
    try:
        from config.settings import config
        
        # Check embeddings
        import numpy as np
        embeddings = np.load(os.path.join(config.VECTOR_STORE_PATH, "embeddings.npy"))
        print(f"✅ Embeddings loaded: {embeddings.shape}")
        
        # Check FAISS
        import faiss
        faiss_index = faiss.read_index(os.path.join(config.VECTOR_STORE_PATH, "faiss_index.bin"))
        print(f"✅ FAISS index loaded: {faiss_index.ntotal} vectors")
        
        # Check chunks
        import pickle
        with open(os.path.join(config.VECTOR_STORE_PATH, "chunks.pkl"), "rb") as f:
            chunks = pickle.load(f)
        print(f"✅ Chunks loaded: {len(chunks)} chunks")
        
        # Check BM25
        with open(os.path.join(config.VECTOR_STORE_PATH, "bm25.pkl"), "rb") as f:
            bm25 = pickle.load(f)
        print(f"✅ BM25 loaded")
        
    except Exception as e:
        print(f"❌ Vector store check failed: {e}")

def check_model_loading():
    """Check if embedding model loads"""
    print("\n=== Model Loading Check ===")
    
    try:
        from config.settings import config
        from sentence_transformers import SentenceTransformer
        
        print("Loading embedding model...")
        model = SentenceTransformer(config.EMBEDDING_MODEL, device='cpu')
        print("✅ Embedding model loaded")
        
        # Test encoding
        test_embedding = model.encode(["test"], show_progress_bar=False)
        print(f"✅ Test encoding successful: shape {test_embedding.shape}")
        
    except Exception as e:
        print(f"❌ Model loading failed: {e}")

def check_api_keys():
    """Check API keys"""
    print("\n=== API Keys Check ===")
    
    openrouter_key = os.getenv("OPENROUTER_API_KEY")
    tavily_key = os.getenv("TAVILY_API_KEY")
    
    if openrouter_key:
        print("✅ OPENROUTER_API_KEY found")
    else:
        print("❌ OPENROUTER_API_KEY missing")
    
    if tavily_key:
        print("✅ TAVILY_API_KEY found")
    else:
        print("⚠️ TAVILY_API_KEY missing (optional)")

if __name__ == "__main__":
    print("🔍 MANIT RAG Chatbot Diagnostic")
    print("=" * 50)
    
    check_files()
    check_imports()
    check_vector_store()
    check_model_loading()
    check_api_keys()
    
    print("\n" + "=" * 50)
    print("Diagnostic complete!")