import os
import sys

# EMERGENCY: Redirect model cache to D: drive (C: is full)
os.environ["HF_HOME"] = r"D:\AI_C\Models\.cache"
os.environ["TRANSFORMERS_CACHE"] = r"D:\AI_C\Models\.cache"

import asyncio
import pandas as pd
import joblib

# Ensure the script can find the app package
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from app.services.ml_pipeline import ml_pipeline
from app.services.db_service import db_service

async def check_databases():
    print("\n--- Step 1: Database Connectivity Check ---")
    mongo_ok = False
    neo4j_ok = False
    
    try:
        # Check MongoDB
        if db_service.mongo_client:
            db_service.mongo_client.admin.command('ping')
            print("DONE: MongoDB: Connected")
            mongo_ok = True
        else:
            print("FAIL: MongoDB: Client not initialized")
    except Exception as e:
        print(f"FAIL: MongoDB: Connection failed ({e})")
        
    try:
        # Check Neo4j
        if db_service.neo4j_driver:
            with db_service.neo4j_driver.session() as session:
                session.run("RETURN 1")
            print("DONE: Neo4j: Connected")
            neo4j_ok = True
        else:
            print("FAIL: Neo4j: Driver not initialized")
    except Exception as e:
        print(f"FAIL: Neo4j: Connection failed ({e})")
        
    return mongo_ok and neo4j_ok

def initialize_models():
    print("\n--- Step 2: Model Initialization (Downloading Pretrained Weights) ---")
    print("This may take several minutes on first run (BERT, RoBERTa, SBERT)...")
    try:
        ml_pipeline._load_models()
        print("DONE: Models: All weights cached/loaded")
    except Exception as e:
        print(f"FAIL: Models: Loading failed ({e})")

def setup_clustering():
    print("\n--- Step 3: Dataset-Specific clustering (10k Sample) ---")
    dataset_path = r"d:\AI_C\DATASETS\Global News Dataset\data.csv"
    if not os.path.exists(dataset_path):
        print(f"WARN: Dataset not found at {dataset_path}. Skipping clustering fit.")
        return

    try:
        print(f"Loading titles from {dataset_path}...")
        df = pd.read_csv(dataset_path, on_bad_lines='skip', nrows=10000)
        titles = df['title'].astype(str).tolist()
        ml_pipeline.refit_clustering(titles, sample_size=10000)
        print("DONE: Clustering: Model calibrated to local dataset")
    except Exception as e:
        print(f"FAIL: Clustering: Fit failed ({e})")

def seed_graph():
    print("\n--- Step 4: GDELT Graph Seeding ---")
    try:
        import subprocess
        print("Running GDELT ingestion...")
        # Use simple python check to avoid env issues
        result = subprocess.run([sys.executable, "ingest_gdelt_graph.py"], 
                                capture_output=True, text=True)
        if result.returncode == 0:
            print("DONE: GDELT Graph: Seeded successfully")
        else:
            print(f"FAIL: GDELT Graph: Ingestion failed\n{result.stderr}")
    except Exception as e:
        print(f"FAIL: GDELT Graph: Script execution failed ({e})")

async def main():
    print("==========================================")
    print("   NARRATIVENET SYSTEM INITIALIZATION    ")
    print("==========================================")
    
    dbs_ready = await check_databases()
    if not dbs_ready:
        print("\n🛑 FATAL: Database services are required. Please ensure Neo4j and MongoDB are running.")
        # We continue anyway to cache models if possible
    
    initialize_models()
    setup_clustering()
    
    if dbs_ready:
        seed_graph()
    else:
        print("\n⚠️ Skipping graph seeding due to missing database connection.")
        
    print("\n✨ Initialization process finished.")
    print("==========================================")

if __name__ == "__main__":
    asyncio.run(main())