"""
Phase 2 Runner — FAISS Indexing, RAG Semantic Retrieval, LLM Explanation, and Scheduling
Author: Antigravity AI
Date: August 2026

This script coordinates:
1. Synthesizing logs and building/saving the local FAISS index
2. Initializing the LLM Explainer and Log Retriever
3. Mocking telemetry anomalies based on classification/RUL alerts
4. Querying the FAISS index for past cases (RAG context)
5. Generating natural language explanations for the production manager
6. Executing scheduling reassignments to mitigate risk
"""

import os
import sys
import logging

# Ensure src/ is in the import path
BASE_DIR = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, os.path.join(BASE_DIR, 'src'))

from utils import set_seeds, ProductionScheduler
from rag import build_and_save_index, LogRetriever, LLMExplainer

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


def main():
    set_seeds(42)
    
    # Define directories according to clean architecture
    data_raw_dir = os.path.join(BASE_DIR, 'data', 'raw')
    data_processed_dir = os.path.join(BASE_DIR, 'data', 'processed')
    models_rag_dir = os.path.join(BASE_DIR, 'models', 'rag')
    reports_dir = os.path.join(BASE_DIR, 'reports')
    
    os.makedirs(data_raw_dir, exist_ok=True)
    os.makedirs(data_processed_dir, exist_ok=True)
    os.makedirs(models_rag_dir, exist_ok=True)
    os.makedirs(reports_dir, exist_ok=True)
    
    # Paths
    sch_path = os.path.join(data_raw_dir, 'production_schedule.csv')
    mach_path = os.path.join(data_raw_dir, 'PdM_machines.csv')
    adj_path = os.path.join(data_processed_dir, 'adjusted_schedule.csv')
    
    logger.info("=========================================")
    logger.info("STARTING PHASE 2: FAISS, LLM, AND OPERATION SCHEDULING")
    logger.info("=========================================")
    
    # ----------------------------------------------------
    # STEP 1: BUILD FAISS INDEX
    # ----------------------------------------------------
    logger.info("\n--- STEP 1: Building local FAISS vector database index ---")
    build_and_save_index(data_raw_dir, models_rag_dir)
    
    # ----------------------------------------------------
    # STEP 2: INITIALIZE RETRIEVER AND EXPLAINER
    # ----------------------------------------------------
    logger.info("\n--- STEP 2: Loading retriever and explainer modules ---")
    retriever = LogRetriever(models_rag_dir)
    
    # To run instantly without waiting to download the 270MB LLM locally,
    # set use_fallback=True. Set to False if you want to download and test SmolLM2-135M.
    use_fallback = True
    if use_fallback:
        logger.info("Using rule-based fallback explainer template for instant local test.")
    else:
        logger.info("Downloading and loading SmolLM2-135M locally...")
        
    explainer = LLMExplainer(force_fallback=use_fallback)
    
    # ----------------------------------------------------
    # STEP 3: SIMULATE ANOMALOUS TELEMETRY & RUN RAG
    # ----------------------------------------------------
    logger.info("\n--- STEP 3: Simulating machine anomaly alerts & executing RAG queries ---")
    
    critical_alerts = {
        3: {'failure_probability': 0.88, 'rul': 12.5},   # Machine 3 high failure risk, low RUL
        12: {'failure_probability': 0.94, 'rul': 8.0}     # Machine 12 high failure risk, low RUL
    }
    
    reports = {}
    for mid, alert in critical_alerts.items():
        query_str = f"Machine ID: {mid} component error telemetry anomalies"
        logger.info(f"Querying FAISS for: '{query_str}'")
        
        # Retrieve top 2 matching historical records
        hits = retriever.query(query_str, k=2)
        
        # Generate LLM explanation using retrieved records as context
        report = explainer.generate_explanation(
            machine_id=mid,
            failure_prob=alert['failure_probability'],
            rul=alert['rul'],
            historical_logs=hits
        )
        reports[mid] = report
        
        # Save report to reports folder
        report_path = os.path.join(reports_dir, f"machine_{mid}_explanation_report.md")
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write(report)
            
        logger.info(f"Report for Machine {mid} saved to {report_path}")
        print("\n" + "="*50)
        print(f"EXPLANATION REPORT FOR MACHINE {mid}:")
        print("="*50)
        print(report)
        print("="*50 + "\n")
        
    # ----------------------------------------------------
    # STEP 4: RUN CLOSED-LOOP PRODUCTION SCHEDULER
    # ----------------------------------------------------
    logger.info("\n--- STEP 4: Running Production Scheduler to adjust operations ---")
    scheduler = ProductionScheduler(sch_path, mach_path)
    
    # Reroute jobs scheduled on at-risk machines (3 and 12) to healthy backups
    adjusted_schedule = scheduler.adjust_schedule(
        predictions=critical_alerts,
        output_adjusted_path=adj_path,
        failure_threshold=0.7,
        min_rul_hours=48.0
    )
    
    logger.info("=========================================")
    logger.info("PHASE 2 COMPLETED SUCCESSFULLY.")
    logger.info("=========================================")


if __name__ == '__main__':
    main()