from __future__ import annotations

import json
import logging
import os
import requests
from typing import Any, Dict, List, Optional

from storage.config import load_config_from_env
from storage.s3_adapter import S3Storage

logger = logging.getLogger(__name__)


def retrieve_content_from_url(url: str, timeout: int = 10) -> Optional[str]:
    """
    Retrieve full content from a URL for LLM analysis
    """
    try:
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
        }
        
        response = requests.get(url, headers=headers, timeout=timeout)
        response.raise_for_status()
        
        # Basic text extraction (you might want to use a proper HTML parser)
        content = response.text
        
        # Simple cleanup - remove HTML tags and extra whitespace
        import re
        content = re.sub(r'<[^>]+>', ' ', content)
        content = re.sub(r'\s+', ' ', content).strip()
        
        return content[:10000]  # Limit to 10k characters for LLM processing
        
    except Exception as e:
        logger.warning(f"Failed to retrieve content from {url}: {e}")
        return None


def retrieve_processed_json_files(s3: S3Storage, bucket: str, limit: Optional[int] = None) -> List[Dict[str, Any]]:
    """
    Retrieve processed JSON files from S3 bucket
    """
    try:
        logger.info(f"Retrieving processed JSON files from bucket: {bucket}")
        
        # List all objects in the processed 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]
        
        logger.info(f"Found {len(json_objects)} JSON files to process")
        
        # Retrieve and parse each JSON file
        processed_files = []
        for obj_key in json_objects:
            try:
                logger.info(f"Processing file: {obj_key}")
                json_content = s3.get_json(bucket, obj_key)
                
                if json_content:
                    # Add metadata about the source file
                    json_content['_source_file'] = obj_key
                    json_content['_bucket'] = bucket
                    processed_files.append(json_content)
                    
            except Exception as e:
                logger.error(f"Failed to process file {obj_key}: {e}")
                continue
        
        logger.info(f"Successfully retrieved {len(processed_files)} processed files")
        return processed_files
        
    except Exception as e:
        logger.error(f"Failed to retrieve processed JSON files: {e}")
        return []


def prepare_content_for_llm(processed_files: List[Dict[str, Any]], retrieve_full_content: bool = True) -> List[Dict[str, Any]]:
    """
    Prepare content from processed files for LLM analysis
    """
    llm_ready_content = []
    
    for file_data in processed_files:
        try:
            # Extract relevant fields for LLM analysis
            content = {
                'url': file_data.get('url', ''),
                'title': file_data.get('title', ''),
                'brief_description': file_data.get('brief_description', ''),
                'author': file_data.get('author', ''),
                'date': file_data.get('date', ''),
                'language': file_data.get('language', ''),
                'relevant_features': file_data.get('relevant_features', []),
                'relevance_score': file_data.get('relevance_score', 0.0),
                'source_file': file_data.get('_source_file', ''),
                'bucket': file_data.get('_bucket', ''),
                'content_retrieval_required': file_data.get('content_retrieval_required', False)
            }
            
            # If content retrieval is required and enabled, fetch full content from URL
            if retrieve_full_content and content['content_retrieval_required'] and content['url']:
                logger.info(f"Retrieving full content from URL: {content['url']}")
                full_content = retrieve_content_from_url(content['url'])
                if full_content:
                    content['extracted_text'] = full_content
                    content['content_retrieved'] = True
                else:
                    content['extracted_text'] = content['brief_description']  # Fallback to brief description
                    content['content_retrieved'] = False
            else:
                # Use brief description as content
                content['extracted_text'] = content['brief_description']
                content['content_retrieved'] = False
            
            # Only include content with substantial information
            if content['extracted_text'] and len(content['extracted_text']) > 50:
                llm_ready_content.append(content)
                
        except Exception as e:
            logger.error(f"Failed to prepare content for LLM: {e}")
            continue
    
    logger.info(f"Prepared {len(llm_ready_content)} items for LLM analysis")
    return llm_ready_content


def create_llm_prompt(content: Dict[str, Any], feature_keywords: List[str]) -> str:
    """
    Create a prompt for LLM to analyze content for specific features
    """
    # Add content retrieval status to the prompt
    content_status = "Full content retrieved from URL" if content.get('content_retrieved', False) else "Brief description only"
    
    prompt = f"""
Analyze the following cybersecurity content for ransomware-related features and indicators.

CONTENT TO ANALYZE:
URL: {content['url']}
Title: {content['title']}
Author: {content['author']}
Date: {content['date']}
Language: {content['language']}
Relevance Score: {content['relevance_score']}
Content Status: {content_status}

TEXT CONTENT:
{content['extracted_text'][:4000]}  # Limit to first 4000 chars for LLM

KEY FEATURES TO LOOK FOR:
{', '.join(feature_keywords)}

Please analyze this content and provide:
1. Ransomware family identification (if any)
2. Attack techniques mentioned
3. Target industries/organizations
4. Technical indicators (file extensions, encryption methods, etc.)
5. Timeline information
6. Threat actor information
7. Impact assessment
8. Confidence level (1-10)
9. Whether full content retrieval would be beneficial for deeper analysis

Format your response as structured JSON.
"""

    return prompt


def batch_process_for_llm(processed_files: List[Dict[str, Any]], batch_size: int = 10) -> List[List[Dict[str, Any]]]:
    """
    Split processed files into batches for LLM processing
    """
    batches = []
    for i in range(0, len(processed_files), batch_size):
        batch = processed_files[i:i + batch_size]
        batches.append(batch)
    
    logger.info(f"Created {len(batches)} batches of size {batch_size}")
    return batches


def lambda_handler(event, context):
    """
    Lambda handler for processing JSON files from S3 for LLM analysis
    """
    try:
        cfg = load_config_from_env()
        s3 = S3Storage(cfg, no_aws=False)
        
        # Get parameters from event
        bucket = event.get('bucket', cfg.s3_bucket_processed)
        limit = event.get('limit', None)
        batch_size = event.get('batch_size', 10)
        
        logger.info(f"Processing JSON files from bucket: {bucket}")
        
        # Retrieve processed JSON files
        processed_files = retrieve_processed_json_files(s3, bucket, limit)
        
        if not processed_files:
            return {
                'statusCode': 200,
                'body': json.dumps({
                    'message': 'No processed JSON files found',
                    'count': 0,
                    'batches': 0
                })
            }
        
        # Prepare content for LLM (with content retrieval enabled)
        retrieve_content = event.get('retrieve_content', True)
        llm_ready_content = prepare_content_for_llm(processed_files, retrieve_full_content=retrieve_content)
        
        # Create batches for processing
        batches = batch_process_for_llm(llm_ready_content, batch_size)
        
        # Load feature keywords for analysis
        feature_keywords = [
            'ransomware', 'lockbit', 'cl0p', 'royal', 'conti', 'maze', 'ryuk',
            'encryption', 'cryptocurrency', 'bitcoin', 'monero', 'ethereum',
            'phishing', 'social-engineering', 'lateral-movement', 'privilege-escalation',
            'command-control', 'c2', 'exfiltration', 'double-extortion',
            'healthcare', 'hospital', 'finance', 'banking', 'government', 'education'
        ]
        
        # Create prompts for each batch
        batch_prompts = []
        for batch in batches:
            batch_prompt = {
                'batch_id': len(batch_prompts),
                'items': len(batch),
                'prompts': []
            }
            
            for item in batch:
                prompt = create_llm_prompt(item, feature_keywords)
                batch_prompt['prompts'].append({
                    'content': item,
                    'prompt': prompt
                })
            
            batch_prompts.append(batch_prompt)
        
        # Store results in S3 for further processing
        results_key = f"llm-analysis/batch-analysis-{context.aws_request_id}.json"
        s3.put_json(bucket, "llm-analysis", context.aws_request_id, {
            'total_files': len(processed_files),
            'llm_ready_files': len(llm_ready_content),
            'batches': len(batches),
            'batch_prompts': batch_prompts,
            'feature_keywords': feature_keywords,
            'processed_at': context.aws_request_id
        })
        
        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Successfully processed files for LLM analysis',
                'total_files': len(processed_files),
                'llm_ready_files': len(llm_ready_content),
                'batches': len(batches),
                'results_s3_key': results_key
            })
        }
        
    except Exception as e:
        logger.error(f"Error processing files for LLM: {e}")
        return {
            'statusCode': 500,
            'body': json.dumps({
                'error': str(e),
                'message': 'Failed to process files for LLM analysis'
            })
        }
