#!/usr/bin/env python3
"""
LLM Feature Matcher for processed JSON content
This script takes the prepared LLM analysis data and processes it through an LLM
"""

import os
import sys
import json
import logging
from typing import Dict, Any, List, 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'
)

class LLMFeatureMatcher:
    """LLM Feature Matcher for cybersecurity content analysis"""
    
    def __init__(self, llm_provider: str = "openai"):
        """
        Initialize the LLM Feature Matcher
        
        Args:
            llm_provider: The LLM provider to use ("openai", "anthropic", "local")
        """
        self.llm_provider = llm_provider
        self.setup_llm_client()
    
    def setup_llm_client(self):
        """Setup the LLM client based on provider"""
        try:
            if self.llm_provider == "openai":
                try:
                    import openai
                    self.client = openai.OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
                    print("✅ OpenAI client initialized")
                except ImportError:
                    print("❌ OpenAI library not installed. Install with: pip install openai")
                    raise
            elif self.llm_provider == "anthropic":
                try:
                    import anthropic
                    self.client = anthropic.Anthropic(api_key=os.getenv('ANTHROPIC_API_KEY'))
                    print("✅ Anthropic client initialized")
                except ImportError:
                    print("❌ Anthropic library not installed. Install with: pip install anthropic")
                    raise
            else:
                print("⚠️  Unknown LLM provider. Using mock client for testing.")
                self.client = None
                
        except Exception as e:
            print(f"❌ Failed to initialize LLM client: {e}")
            self.client = None
    
    def analyze_with_llm(self, prompt: str) -> Dict[str, Any]:
        """
        Send prompt to LLM and get analysis results
        
        Args:
            prompt: The prompt to send to the LLM
            
        Returns:
            Dictionary containing the LLM response
        """
        try:
            if not self.client:
                return self.mock_llm_response()
            
            if self.llm_provider == "openai":
                return self.analyze_with_openai(prompt)
            elif self.llm_provider == "anthropic":
                return self.analyze_with_anthropic(prompt)
            else:
                return self.mock_llm_response()
                
        except Exception as e:
            print(f"❌ Error calling LLM: {e}")
            return {
                "error": str(e),
                "analysis": None
            }
    
    def analyze_with_openai(self, prompt: str) -> Dict[str, Any]:
        """Analyze using OpenAI API"""
        try:
            response = self.client.chat.completions.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": "You are a cybersecurity analyst specializing in ransomware threat intelligence. Always respond with valid JSON."},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.1,
                max_tokens=4000
            )
            
            content = response.choices[0].message.content
            
            # Try to parse as JSON
            try:
                parsed_content = json.loads(content)
                return {
                    "success": True,
                    "analysis": parsed_content,
                    "raw_response": content
                }
            except json.JSONDecodeError:
                return {
                    "success": False,
                    "analysis": None,
                    "raw_response": content,
                    "error": "Response is not valid JSON"
                }
                
        except Exception as e:
            return {
                "success": False,
                "analysis": None,
                "error": str(e)
            }
    
    def analyze_with_anthropic(self, prompt: str) -> Dict[str, Any]:
        """Analyze using Anthropic API"""
        try:
            response = self.client.messages.create(
                model="claude-3-sonnet-20240229",
                max_tokens=4000,
                temperature=0.1,
                messages=[{"role": "user", "content": prompt}]
            )
            
            content = response.content[0].text
            
            # Try to parse as JSON
            try:
                parsed_content = json.loads(content)
                return {
                    "success": True,
                    "analysis": parsed_content,
                    "raw_response": content
                }
            except json.JSONDecodeError:
                return {
                    "success": False,
                    "analysis": None,
                    "raw_response": content,
                    "error": "Response is not valid JSON"
                }
                
        except Exception as e:
            return {
                "success": False,
                "analysis": None,
                "error": str(e)
            }
    
    def mock_llm_response(self) -> Dict[str, Any]:
        """Mock LLM response for testing without API keys"""
        return {
            "success": True,
            "analysis": {
                "batch_analysis": {
                    "total_documents": 1,
                    "ransomware_families_found": ["MockRansomware"],
                    "attack_techniques": ["Phishing", "Encryption"],
                    "target_industries": ["Healthcare"],
                    "technical_indicators": [".locked", "bitcoin"],
                    "timeline_events": [],
                    "threat_actors": ["MockActor"],
                    "confidence_level": 7,
                    "key_findings": ["Mock analysis for testing"]
                },
                "document_analyses": [
                    {
                        "document_id": 1,
                        "url": "https://example.com",
                        "ransomware_family": "MockRansomware",
                        "attack_techniques": ["Phishing"],
                        "target_industry": "Healthcare",
                        "technical_indicators": [".locked"],
                        "threat_actor": "MockActor",
                        "confidence": 7,
                        "key_insights": ["Mock insight"]
                    }
                ]
            },
            "raw_response": "Mock response for testing",
            "mock": True
        }
    
    def process_batch(self, batch_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process a single batch through the LLM
        
        Args:
            batch_data: Dictionary containing batch information and LLM prompt
            
        Returns:
            Dictionary containing the analysis results
        """
        print(f"🤖 Processing batch {batch_data['batch_id']} with {batch_data['batch_size']} items")
        
        # Get the LLM prompt
        prompt = batch_data['llm_prompt']
        
        # Send to LLM
        llm_response = self.analyze_with_llm(prompt)
        
        # Add metadata
        result = {
            "batch_id": batch_data['batch_id'],
            "processed_at": datetime.now().isoformat(),
            "batch_size": batch_data['batch_size'],
            "llm_provider": self.llm_provider,
            "llm_response": llm_response,
            "source_items": batch_data['items']
        }
        
        return result
    
    def process_all_batches(self, analysis_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process all batches in the analysis data
        
        Args:
            analysis_data: Dictionary containing all batches and metadata
            
        Returns:
            Dictionary containing all analysis results
        """
        print(f"🚀 Processing {len(analysis_data['batches'])} batches through LLM")
        
        results = {
            "metadata": analysis_data['metadata'],
            "processed_at": datetime.now().isoformat(),
            "llm_provider": self.llm_provider,
            "batch_results": [],
            "summary": {
                "total_batches": len(analysis_data['batches']),
                "successful_batches": 0,
                "failed_batches": 0,
                "total_documents_analyzed": 0
            }
        }
        
        for batch in analysis_data['batches']:
            try:
                batch_result = self.process_batch(batch)
                results['batch_results'].append(batch_result)
                
                # Update summary
                if batch_result['llm_response'].get('success', False):
                    results['summary']['successful_batches'] += 1
                    results['summary']['total_documents_analyzed'] += batch_result['batch_size']
                else:
                    results['summary']['failed_batches'] += 1
                    
            except Exception as e:
                print(f"❌ Failed to process batch {batch['batch_id']}: {e}")
                results['summary']['failed_batches'] += 1
                results['batch_results'].append({
                    "batch_id": batch['batch_id'],
                    "error": str(e),
                    "llm_response": {"success": False, "error": str(e)}
                })
        
        return results

def main():
    """Main function to run LLM feature matching"""
    import argparse
    
    parser = argparse.ArgumentParser(description='Run LLM feature matching on processed JSON files')
    parser.add_argument('--input-file', required=True, help='Path to the LLM analysis data JSON file')
    parser.add_argument('--provider', choices=['openai', 'anthropic', 'mock'], default='mock', help='LLM provider to use')
    parser.add_argument('--output-file', help='Output file for results (default: auto-generated)')
    
    args = parser.parse_args()
    
    print("🤖 LLM Feature Matcher")
    print("=" * 50)
    print(f"📁 Input file: {args.input_file}")
    print(f"🤖 LLM Provider: {args.provider}")
    
    # Load the analysis data
    try:
        with open(args.input_file, 'r', encoding='utf-8') as f:
            analysis_data = json.load(f)
        print(f"✅ Loaded analysis data with {len(analysis_data['batches'])} batches")
    except Exception as e:
        print(f"❌ Failed to load input file: {e}")
        return
    
    # Initialize LLM matcher
    matcher = LLMFeatureMatcher(args.provider)
    
    # Process all batches
    results = matcher.process_all_batches(analysis_data)
    
    # Generate output filename
    if not args.output_file:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        args.output_file = f"llm_analysis_results_{timestamp}.json"
    
    # Save results
    with open(args.output_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    
    print(f"\n✅ LLM Analysis Complete!")
    print(f"📊 Results:")
    print(f"  Total batches: {results['summary']['total_batches']}")
    print(f"  Successful: {results['summary']['successful_batches']}")
    print(f"  Failed: {results['summary']['failed_batches']}")
    print(f"  Documents analyzed: {results['summary']['total_documents_analyzed']}")
    print(f"📁 Results saved to: {args.output_file}")

if __name__ == "__main__":
    main()
