#!/usr/bin/env python3
"""
Complete LLM Analysis Pipeline
This script retrieves JSON files from S3, prepares them for LLM analysis, and processes them
"""

import os
import sys
import json
import logging
from typing import Dict, Any, List
from datetime import datetime

# Add the project root to the path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# Set up logging
logging.basicConfig(
    level=logging.INFO, 
    format='%(asctime)s [%(name)s] %(levelname)s: %(message)s'
)

def run_complete_llm_analysis(
    bucket_name: str = None,
    limit: int = 10,
    batch_size: int = 3,
    llm_provider: str = "mock",
    retrieve_content: bool = True
) -> Dict[str, Any]:
    """
    Run the complete LLM analysis pipeline
    
    Args:
        bucket_name: S3 bucket name (defaults to processed bucket)
        limit: Maximum number of files to process
        batch_size: Number of documents per batch
        llm_provider: LLM provider to use
        retrieve_content: Whether to retrieve full content from URLs
        
    Returns:
        Dictionary containing all analysis results
    """
    print("🚀 Complete LLM Analysis Pipeline")
    print("=" * 60)
    
    # Step 1: Retrieve JSON files from S3
    print("📁 Step 1: Retrieving JSON files from S3...")
    from process_s3_json_for_llm import retrieve_and_analyze_s3_json, prepare_llm_batch, save_llm_analysis_results
    
    result = retrieve_and_analyze_s3_json(bucket_name, limit)
    files = result.get('files', [])
    analysis_summary = result.get('analysis', {})
    
    if not files:
        print("❌ No files found to process")
        return {"error": "No files found"}
    
    print(f"✅ Retrieved {len(files)} files from S3")
    
    # Step 2: Prepare batches for LLM
    print("📦 Step 2: Preparing LLM batches...")
    batches = prepare_llm_batch(files, batch_size)
    
    if not batches:
        print("❌ No valid batches prepared")
        return {"error": "No valid batches"}
    
    print(f"✅ Prepared {len(batches)} batches for LLM analysis")
    
    # Step 3: Create analysis data package
    print("📋 Step 3: Creating analysis data package...")
    analysis_data = {
        'metadata': {
            'created_at': datetime.now().isoformat(),
            'total_files': analysis_summary.get('total_files', 0),
            'successful_retrievals': analysis_summary.get('successful_retrievals', 0),
            'batches_prepared': len(batches),
            'analysis_summary': analysis_summary,
            'bucket_name': bucket_name,
            'limit': limit,
            'batch_size': batch_size
        },
        'batches': []
    }
    
    # Add each batch with its LLM prompt
    for batch in batches:
        from process_s3_json_for_llm import create_llm_analysis_prompt
        llm_prompt = create_llm_analysis_prompt(batch)
        
        analysis_data['batches'].append({
            'batch_id': batch['batch_id'],
            'items': batch['items'],
            'llm_prompt': llm_prompt,
            'batch_size': batch['batch_size']
        })
    
    print(f"✅ Created analysis data package")
    
    # Step 4: Process through LLM
    print(f"🤖 Step 4: Processing through LLM ({llm_provider})...")
    from llm_feature_matcher import LLMFeatureMatcher
    
    matcher = LLMFeatureMatcher(llm_provider)
    llm_results = matcher.process_all_batches(analysis_data)
    
    print(f"✅ LLM analysis complete")
    
    # Step 5: Combine results
    print("📊 Step 5: Combining results...")
    final_results = {
        'pipeline_metadata': {
            'created_at': datetime.now().isoformat(),
            'bucket_name': bucket_name,
            'limit': limit,
            'batch_size': batch_size,
            'llm_provider': llm_provider
        },
        's3_analysis': analysis_summary,
        'llm_results': llm_results,
        'summary': {
            'total_files_retrieved': len(files),
            'batches_processed': len(batches),
            'successful_llm_batches': llm_results['summary']['successful_batches'],
            'failed_llm_batches': llm_results['summary']['failed_batches'],
            'total_documents_analyzed': llm_results['summary']['total_documents_analyzed']
        }
    }
    
    print(f"✅ Results combined")
    
    return final_results

def save_results(results: Dict[str, Any], filename: str = None) -> str:
    """Save results to JSON file"""
    if not filename:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"complete_llm_analysis_{timestamp}.json"
    
    with open(filename, 'w', encoding='utf-8') as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    
    print(f"💾 Results saved to: {filename}")
    return filename

def print_summary(results: Dict[str, Any]):
    """Print a summary of the analysis results"""
    print("\n" + "=" * 60)
    print("📊 ANALYSIS SUMMARY")
    print("=" * 60)
    
    # S3 Analysis Summary
    s3_analysis = results.get('s3_analysis', {})
    print(f"📁 S3 Analysis:")
    print(f"  Total files found: {s3_analysis.get('total_files', 0)}")
    print(f"  Successfully retrieved: {s3_analysis.get('successful_retrievals', 0)}")
    print(f"  High relevance files: {s3_analysis.get('high_relevance_files', 0)}")
    print(f"  Files with features: {s3_analysis.get('files_with_features', 0)}")
    
    # LLM Analysis Summary
    llm_summary = results.get('llm_results', {}).get('summary', {})
    print(f"\n🤖 LLM Analysis:")
    print(f"  Total batches: {llm_summary.get('total_batches', 0)}")
    print(f"  Successful batches: {llm_summary.get('successful_batches', 0)}")
    print(f"  Failed batches: {llm_summary.get('failed_batches', 0)}")
    print(f"  Documents analyzed: {llm_summary.get('total_documents_analyzed', 0)}")
    
    # Pipeline Summary
    pipeline_summary = results.get('summary', {})
    print(f"\n🚀 Pipeline Summary:")
    print(f"  Files retrieved: {pipeline_summary.get('total_files_retrieved', 0)}")
    print(f"  Batches processed: {pipeline_summary.get('batches_processed', 0)}")
    print(f"  LLM success rate: {pipeline_summary.get('successful_llm_batches', 0)}/{pipeline_summary.get('batches_processed', 1)}")
    
    # Feature Analysis
    feature_counts = s3_analysis.get('feature_counts', {})
    if feature_counts:
        print(f"\n🎯 Top Features Found:")
        top_features = sorted(feature_counts.items(), key=lambda x: x[1], reverse=True)[:10]
        for feature, count in top_features:
            print(f"  {feature}: {count}")

def main():
    """Main function"""
    import argparse
    
    parser = argparse.ArgumentParser(description='Run complete LLM analysis pipeline')
    parser.add_argument('--bucket', help='S3 bucket name (defaults to processed bucket)')
    parser.add_argument('--limit', type=int, default=10, help='Maximum files to process (default: 10)')
    parser.add_argument('--batch-size', type=int, default=3, help='Documents per batch (default: 3)')
    parser.add_argument('--provider', choices=['openai', 'anthropic', 'mock'], default='mock', help='LLM provider (default: mock)')
    parser.add_argument('--retrieve-content', type=str, default='true', help='Retrieve full content from URLs (default: true)')
    parser.add_argument('--output', help='Output filename (default: auto-generated)')
    
    args = parser.parse_args()
    
    print(f"🔧 Configuration:")
    print(f"  Bucket: {args.bucket or 'default processed bucket'}")
    print(f"  Limit: {args.limit} files")
    print(f"  Batch size: {args.batch_size}")
    print(f"  LLM Provider: {args.provider}")
    print()
    
    # Run the complete pipeline
    try:
        results = run_complete_llm_analysis(
            bucket_name=args.bucket,
            limit=args.limit,
            batch_size=args.batch_size,
            llm_provider=args.provider,
            retrieve_content=args.retrieve_content.lower() == 'true'
        )
        
        if 'error' in results:
            print(f"❌ Pipeline failed: {results['error']}")
            return
        
        # Save results
        output_file = save_results(results, args.output)
        
        # Print summary
        print_summary(results)
        
        print(f"\n🎉 Complete LLM Analysis Pipeline Finished!")
        print(f"📁 Results saved to: {output_file}")
        
    except Exception as e:
        print(f"❌ Pipeline failed: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()
