#!/usr/bin/env python3
"""
Process JSON files from S3 for LLM analysis
This script retrieves processed JSON files from S3 and prepares them for LLM feature matching
"""

import os
import sys
import json
import logging
from typing import List, Dict, Any, Optional
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 retrieve_and_analyze_s3_json(bucket_name: str = None, limit: int = None) -> Dict[str, Any]:
    """
    Retrieve and analyze JSON files from S3 for LLM processing
    """
    try:
        from storage.config import load_config_from_env
        from storage.s3_adapter import S3Storage
        
        # Load configuration
        cfg = load_config_from_env()
        s3 = S3Storage(cfg, no_aws=False)
        
        # Use provided bucket or default processed bucket
        bucket = bucket_name or cfg.s3_bucket_processed
        
        print(f"🔍 Retrieving JSON files from S3 bucket: {bucket}")
        
        # List all objects in the bucket
        objects = s3.list_objects(bucket)
        
        # Filter for JSON files in the processed folder
        json_objects = [obj for obj in objects if obj.endswith('.json') and 'processed/' in obj]
        
        if limit:
            json_objects = json_objects[:limit]
        
        print(f"📁 Found {len(json_objects)} JSON files to analyze")
        
        # Retrieve and analyze each JSON file
        processed_files = []
        analysis_summary = {
            'total_files': len(json_objects),
            'successful_retrievals': 0,
            'failed_retrievals': 0,
            'high_relevance_files': 0,
            'files_with_features': 0,
            'total_content_length': 0,
            'languages': {},
            'domains': {},
            'feature_counts': {}
        }
        
        for obj_key in json_objects:
            try:
                print(f"📄 Processing: {obj_key}")
                json_content = s3.get_json(bucket, obj_key)
                
                if json_content:
                    # Add metadata
                    json_content['_source_file'] = obj_key
                    json_content['_bucket'] = bucket
                    processed_files.append(json_content)
                    
                    # Update analysis summary
                    analysis_summary['successful_retrievals'] += 1
                    
                    # Count high relevance files
                    relevance_score = json_content.get('relevance_score', 0)
                    if relevance_score > 0.5:
                        analysis_summary['high_relevance_files'] += 1
                    
                    # Count files with features
                    relevant_features = json_content.get('relevant_features', [])
                    if relevant_features:
                        analysis_summary['files_with_features'] += 1
                        
                        # Count individual features
                        for feature in relevant_features:
                            analysis_summary['feature_counts'][feature] = analysis_summary['feature_counts'].get(feature, 0) + 1
                    
                    # Track content length
                    text_length = len(json_content.get('extracted_text', ''))
                    analysis_summary['total_content_length'] += text_length
                    
                    # Track languages
                    language = json_content.get('language', 'unknown')
                    analysis_summary['languages'][language] = analysis_summary['languages'].get(language, 0) + 1
                    
                    # Track domains
                    url = json_content.get('url', '')
                    if url:
                        try:
                            from urllib.parse import urlparse
                            domain = urlparse(url).netloc
                            analysis_summary['domains'][domain] = analysis_summary['domains'].get(domain, 0) + 1
                        except:
                            pass
                    
            except Exception as e:
                print(f"❌ Failed to process {obj_key}: {e}")
                analysis_summary['failed_retrievals'] += 1
                continue
        
        print(f"✅ Successfully processed {analysis_summary['successful_retrievals']} files")
        
        return {
            'files': processed_files,
            'analysis': analysis_summary
        }
        
    except Exception as e:
        print(f"❌ Error retrieving files from S3: {e}")
        return {'files': [], 'analysis': {}}

def prepare_llm_batch(files: List[Dict[str, Any]], batch_size: int = 5) -> List[Dict[str, Any]]:
    """
    Prepare files in batches for LLM processing
    """
    print(f"📦 Preparing {len(files)} files in batches of {batch_size}")
    
    batches = []
    for i in range(0, len(files), batch_size):
        batch = files[i:i + batch_size]
        
        # Prepare each item in the batch for LLM
        llm_batch = []
        for item in batch:
            llm_item = {
                'url': item.get('url', ''),
                'title': item.get('title', ''),
                'extracted_text': item.get('extracted_text', '')[:3000],  # Limit for LLM
                'author': item.get('author', ''),
                'date': item.get('date', ''),
                'language': item.get('language', ''),
                'relevant_features': item.get('relevant_features', []),
                'relevance_score': item.get('relevance_score', 0.0),
                'source_file': item.get('_source_file', ''),
                'content_length': len(item.get('extracted_text', ''))
            }
            
            # Only include items with substantial content
            if llm_item['extracted_text'] and len(llm_item['extracted_text']) > 200:
                llm_batch.append(llm_item)
        
        if llm_batch:
            batches.append({
                'batch_id': len(batches),
                'items': llm_batch,
                'batch_size': len(llm_batch)
            })
    
    print(f"📦 Created {len(batches)} batches for LLM processing")
    return batches

def create_llm_analysis_prompt(batch: Dict[str, Any]) -> str:
    """
    Create a comprehensive prompt for LLM analysis of a batch of files
    """
    prompt = f"""
You are a cybersecurity analyst specializing in ransomware threat intelligence. 
Analyze the following batch of {batch['batch_size']} documents for ransomware-related features and indicators.

BATCH ANALYSIS REQUEST:
- Identify ransomware families mentioned
- Extract attack techniques and TTPs
- Identify target industries/organizations
- Extract technical indicators (file extensions, encryption methods, etc.)
- Identify timeline information
- Extract threat actor information
- Assess impact and confidence levels

DOCUMENTS TO ANALYZE:
"""

    for i, item in enumerate(batch['items'], 1):
        prompt += f"""
DOCUMENT {i}:
URL: {item['url']}
Title: {item['title']}
Author: {item['author']}
Date: {item['date']}
Language: {item['language']}
Relevance Score: {item['relevance_score']}
Content Length: {item['content_length']} characters
Existing Features: {', '.join(item['relevant_features'])}

TEXT CONTENT:
{item['extracted_text']}

---
"""

    prompt += """
Please provide a comprehensive analysis in the following JSON format:
{
  "batch_analysis": {
    "total_documents": number,
    "ransomware_families_found": [list of families],
    "attack_techniques": [list of techniques],
    "target_industries": [list of industries],
    "technical_indicators": [list of indicators],
    "timeline_events": [list of events with dates],
    "threat_actors": [list of actors],
    "confidence_level": number (1-10),
    "key_findings": [list of key findings]
  },
  "document_analyses": [
    {
      "document_id": number,
      "url": "string",
      "ransomware_family": "string or null",
      "attack_techniques": [list],
      "target_industry": "string or null",
      "technical_indicators": [list],
      "threat_actor": "string or null",
      "confidence": number (1-10),
      "key_insights": [list]
    }
  ]
}
"""

    return prompt

def save_llm_analysis_results(batches: List[Dict[str, Any]], analysis_summary: Dict[str, Any]) -> str:
    """
    Save the prepared LLM analysis data to a local file
    """
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"llm_analysis_data_{timestamp}.json"
    
    # Prepare the complete analysis package
    analysis_package = {
        'metadata': {
            'created_at': timestamp,
            'total_files': analysis_summary.get('total_files', 0),
            'successful_retrievals': analysis_summary.get('successful_retrievals', 0),
            'batches_prepared': len(batches),
            'analysis_summary': analysis_summary
        },
        'batches': []
    }
    
    # Add each batch with its LLM prompt
    for batch in batches:
        llm_prompt = create_llm_analysis_prompt(batch)
        
        analysis_package['batches'].append({
            'batch_id': batch['batch_id'],
            'items': batch['items'],
            'llm_prompt': llm_prompt,
            'batch_size': batch['batch_size']
        })
    
    # Save to file
    with open(filename, 'w', encoding='utf-8') as f:
        json.dump(analysis_package, f, indent=2, ensure_ascii=False)
    
    print(f"💾 Saved LLM analysis data to: {filename}")
    return filename

def main():
    """Main function to process S3 JSON files for LLM analysis"""
    print("🤖 S3 JSON to LLM Analysis Processor")
    print("=" * 50)
    
    # Configuration
    bucket_name = os.getenv('S3_BUCKET_PROCESSED', 'ransom-analysis-processed')
    limit = int(os.getenv('LLM_ANALYSIS_LIMIT', '20'))  # Limit files for testing
    batch_size = int(os.getenv('LLM_BATCH_SIZE', '5'))
    
    print(f"📊 Configuration:")
    print(f"  Bucket: {bucket_name}")
    print(f"  Limit: {limit} files")
    print(f"  Batch size: {batch_size}")
    print()
    
    # Retrieve and analyze files
    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
    
    # Print analysis summary
    print("\n📊 ANALYSIS SUMMARY")
    print("=" * 50)
    print(f"Total files found: {analysis_summary.get('total_files', 0)}")
    print(f"Successfully retrieved: {analysis_summary.get('successful_retrievals', 0)}")
    print(f"High relevance files: {analysis_summary.get('high_relevance_files', 0)}")
    print(f"Files with features: {analysis_summary.get('files_with_features', 0)}")
    print(f"Total content length: {analysis_summary.get('total_content_length', 0):,} characters")
    
    print(f"\n🌍 Languages:")
    for lang, count in analysis_summary.get('languages', {}).items():
        print(f"  {lang}: {count}")
    
    print(f"\n🔗 Top domains:")
    domains = sorted(analysis_summary.get('domains', {}).items(), key=lambda x: x[1], reverse=True)
    for domain, count in domains[:10]:
        print(f"  {domain}: {count}")
    
    print(f"\n🎯 Feature counts:")
    features = sorted(analysis_summary.get('feature_counts', {}).items(), key=lambda x: x[1], reverse=True)
    for feature, count in features[:10]:
        print(f"  {feature}: {count}")
    
    # Prepare batches for LLM
    print(f"\n🤖 Preparing LLM analysis batches...")
    batches = prepare_llm_batch(files, batch_size)
    
    # Save results
    filename = save_llm_analysis_results(batches, analysis_summary)
    
    print(f"\n✅ LLM Analysis Preparation Complete!")
    print(f"📁 Results saved to: {filename}")
    print(f"📦 {len(batches)} batches prepared for LLM analysis")
    print(f"📄 Total documents ready for analysis: {sum(batch['batch_size'] for batch in batches)}")
    
    print(f"\n🎯 Next steps:")
    print(f"  1. Use the generated prompts with your preferred LLM")
    print(f"  2. Process each batch through the LLM")
    print(f"  3. Parse the JSON responses for threat intelligence")
    print(f"  4. Store results back to S3 or database")

if __name__ == "__main__":
    main()
