#!/usr/bin/env python3
"""
Script to combine all data for a given day into one file for LLM processing.
This script reads all JSON files from a specific date in the S3 bucket and combines them.
The data format now includes lightweight entries with URLs and brief descriptions,
with content_retrieval_required flags for LLM content retrieval when needed.

Outputs multiple formats:
- JSON: Structured data for programmatic analysis
- Markdown: LLM-friendly narrative format
- CSV: Spreadsheet-compatible export
- ZIP: Combined archive with all formats
"""

import os
import sys
import json
import logging
import csv
import zipfile
from datetime import datetime, timezone
from typing import List, Dict, Any, Optional
from io import StringIO, BytesIO
import boto3
from botocore.exceptions import ClientError

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class DailyDataCombiner:
    def __init__(self, source_bucket: str, combined_bucket: str, region: str = "us-east-1"):
        self.source_bucket = source_bucket
        self.combined_bucket = combined_bucket
        self.s3_client = boto3.client('s3', region_name=region)
        self.s3_resource = boto3.resource('s3', region_name=region)
        
        # Ensure combined bucket exists
        self._ensure_combined_bucket()
    
    def _ensure_combined_bucket(self):
        """Ensure the combined data bucket exists"""
        try:
            self.s3_client.head_bucket(Bucket=self.combined_bucket)
            logger.info(f"Combined bucket exists: {self.combined_bucket}")
        except ClientError as e:
            error_code = e.response['Error']['Code']
            if error_code == '404':
                logger.info(f"Creating combined bucket: {self.combined_bucket}")
                try:
                    self.s3_client.create_bucket(Bucket=self.combined_bucket)
                    logger.info(f"Successfully created bucket: {self.combined_bucket}")
                except ClientError as create_error:
                    logger.warning(f"Could not create bucket: {create_error}")
            else:
                logger.warning(f"Error checking bucket: {e}")
        
    def list_daily_files(self, date_str: str) -> List[str]:
        """List all files for a specific date (YYYY/MM/DD format)"""
        try:
            prefix = f"{date_str}/"
            response = self.s3_client.list_objects_v2(
                Bucket=self.source_bucket,
                Prefix=prefix
            )
            
            files = []
            if 'Contents' in response:
                for obj in response['Contents']:
                    if obj['Key'].endswith('.json'):
                        files.append(obj['Key'])
            
            logger.info(f"Found {len(files)} JSON files for date {date_str}")
            return files
            
        except ClientError as e:
            logger.error(f"Error listing files for date {date_str}: {e}")
            return []
    
    def read_json_file(self, key: str) -> Optional[Dict[str, Any]]:
        """Read and parse a JSON file from S3"""
        try:
            response = self.s3_client.get_object(Bucket=self.source_bucket, Key=key)
            content = response['Body'].read().decode('utf-8')
            return json.loads(content)
        except (ClientError, json.JSONDecodeError) as e:
            logger.warning(f"Error reading file {key}: {e}")
            return None
    
    def combine_daily_data(self, date_str: str) -> Dict[str, Any]:
        """Combine all data for a specific date"""
        files = self.list_daily_files(date_str)
        
        if not files:
            logger.warning(f"No files found for date {date_str}")
            return {
                "date": date_str,
                "total_files": 0,
                "combined_data": [],
                "summary": {
                    "total_articles": 0,
                    "websites": set(),
                    "features_found": set(),
                    "total_relevance_score": 0.0
                }
            }
        
        combined_data = []
        summary = {
            "total_articles": 0,
            "websites": set(),
            "features_found": set(),
            "total_relevance_score": 0.0,
            "languages": set(),
            "error_files": 0
        }
        
        for file_key in files:
            data = self.read_json_file(file_key)
            if data:
                # Extract website from filename
                filename = file_key.split('/')[-1]
                website = filename.split('_')[0] if '_' in filename else "unknown"
                summary["websites"].add(website)
                
                # Add to combined data
                combined_data.append(data)
                summary["total_articles"] += 1
                
                # Extract features
                if 'features' in data:
                    for feature, value in data['features'].items():
                        if value:
                            summary["features_found"].add(feature)
                
                # Add relevance score
                if 'relevance_score' in data:
                    summary["total_relevance_score"] += data['relevance_score']
                
                # Add language
                if 'language' in data:
                    summary["languages"].add(data['language'])
            else:
                summary["error_files"] += 1
        
        # Convert sets to lists for JSON serialization
        summary["websites"] = list(summary["websites"])
        summary["features_found"] = list(summary["features_found"])
        summary["languages"] = list(summary["languages"])
        
        # Calculate average relevance score
        if summary["total_articles"] > 0:
            summary["average_relevance_score"] = summary["total_relevance_score"] / summary["total_articles"]
        else:
            summary["average_relevance_score"] = 0.0
        
        result = {
            "date": date_str,
            "total_files": len(files),
            "combined_data": combined_data,
            "summary": summary,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "created_by": "daily_data_combiner"
        }
        
        logger.info(f"Combined {len(combined_data)} articles from {len(summary['websites'])} websites")
        logger.info(f"Features found: {summary['features_found']}")
        logger.info(f"Average relevance score: {summary['average_relevance_score']:.3f}")
        
        return result
    
    def create_markdown_summary(self, combined_data: Dict[str, Any]) -> str:
        """Create a markdown-formatted summary for LLM analysis"""
        md = []
        
        # Header with metadata
        md.append(f"# Ransomware Intelligence Report - {combined_data['date']}")
        md.append(f"\n**Generated:** {combined_data['created_at']}")
        md.append(f"\n## Summary Statistics\n")
        md.append(f"- **Total Articles:** {combined_data['summary']['total_articles']}")
        md.append(f"- **Unique Websites:** {len(combined_data['summary']['websites'])}")
        md.append(f"- **Average Relevance Score:** {combined_data['summary']['average_relevance_score']:.2f}")
        md.append(f"- **Languages:** {', '.join(combined_data['summary']['languages'])}")
        md.append(f"- **Features Found:** {', '.join(combined_data['summary']['features_found'])}\n")
        
        # Website breakdown
        md.append(f"\n## Sources\n")
        for website in sorted(combined_data['summary']['websites']):
            count = sum(1 for article in combined_data['combined_data'] 
                       if website in article.get('url', '').lower())
            md.append(f"- {website}: {count} articles")
        
        # Articles section
        md.append(f"\n## Articles\n")
        for i, article in enumerate(combined_data['combined_data'], 1):
            md.append(f"\n### Article {i}: {article.get('title', 'No Title')}\n")
            md.append(f"- **URL:** {article.get('url', 'N/A')}")
            md.append(f"- **Author:** {article.get('author', 'Unknown')}")
            md.append(f"- **Date:** {article.get('date', 'N/A')}")
            md.append(f"- **Language:** {article.get('language', 'N/A')}")
            md.append(f"- **Relevance Score:** {article.get('relevance_score', 0.0):.2f}")
            
            if article.get('relevant_features'):
                md.append(f"- **Detected Features:** {', '.join(article['relevant_features'])}")
            
            if article.get('brief_description'):
                md.append(f"\n**Description:**\n{article['brief_description']}\n")
            
            md.append("---")
        
        return '\n'.join(md)
    
    def create_csv_export(self, combined_data: Dict[str, Any]) -> str:
        """Create a CSV export for spreadsheet analysis"""
        output = StringIO()
        writer = csv.writer(output)
        
        # Header
        writer.writerow([
            'Title', 'URL', 'Author', 'Date', 'Language', 
            'Relevance Score', 'Features', 'Description'
        ])
        
        # Data rows
        for article in combined_data['combined_data']:
            writer.writerow([
                article.get('title', ''),
                article.get('url', ''),
                article.get('author', ''),
                article.get('date', ''),
                article.get('language', ''),
                article.get('relevance_score', 0.0),
                '; '.join(article.get('relevant_features', [])),
                article.get('brief_description', '')[:500]  # Truncate long descriptions
            ])
        
        return output.getvalue()
    
    def save_combined_data(self, combined_data: Dict[str, Any], output_key: str) -> str:
        """Save combined data to S3"""
        try:
            # Create a more readable format for LLM processing
            llm_formatted_data = {
                "metadata": {
                    "date": combined_data["date"],
                    "total_articles": combined_data["summary"]["total_articles"],
                    "websites": combined_data["summary"]["websites"],
                    "features_found": combined_data["summary"]["features_found"],
                    "average_relevance_score": combined_data["summary"]["average_relevance_score"],
                    "created_at": combined_data["created_at"]
                },
                "articles": []
            }
            
            # Format articles for LLM consumption
            for article in combined_data["combined_data"]:
                formatted_article = {
                    "url": article.get("url", ""),
                    "title": article.get("title", ""),
                    "author": article.get("author", ""),
                    "date": article.get("date", ""),
                    "language": article.get("language", ""),
                    "brief_description": article.get("brief_description", ""),
                    "content_retrieval_required": article.get("content_retrieval_required", False),
                    "features": article.get("features", {}),
                    "relevant_features": article.get("relevant_features", []),
                    "relevance_score": article.get("relevance_score", 0.0)
                }
                llm_formatted_data["articles"].append(formatted_article)
            
            # Save to S3
            json_content = json.dumps(llm_formatted_data, indent=2, ensure_ascii=False)
            self.s3_client.put_object(
                Bucket=self.combined_bucket,
                Key=output_key,
                Body=json_content.encode('utf-8'),
                ContentType='application/json'
            )
            
            s3_path = f"s3://{self.combined_bucket}/{output_key}"
            logger.info(f"Combined data saved to {s3_path}")
            return s3_path
            
        except ClientError as e:
            logger.error(f"Error saving combined data: {e}")
            raise
    
    def save_combined_data_multi_format(self, combined_data: Dict[str, Any], date_str: str, 
                                         create_zip: bool = True) -> Dict[str, str]:
        """Save combined data in multiple formats (JSON, Markdown, CSV, ZIP)"""
        paths = {}
        date_slug = date_str.replace('/', '-')
        
        try:
            # 1. Create LLM-formatted data structure
            llm_formatted_data = {
                "metadata": {
                    "date": combined_data["date"],
                    "total_articles": combined_data["summary"]["total_articles"],
                    "websites": combined_data["summary"]["websites"],
                    "features_found": combined_data["summary"]["features_found"],
                    "average_relevance_score": combined_data["summary"]["average_relevance_score"],
                    "created_at": combined_data["created_at"]
                },
                "articles": []
            }
            
            for article in combined_data["combined_data"]:
                formatted_article = {
                    "url": article.get("url", ""),
                    "title": article.get("title", ""),
                    "author": article.get("author", ""),
                    "date": article.get("date", ""),
                    "language": article.get("language", ""),
                    "brief_description": article.get("brief_description", ""),
                    "content_retrieval_required": article.get("content_retrieval_required", False),
                    "features": article.get("features", {}),
                    "relevant_features": article.get("relevant_features", []),
                    "relevance_score": article.get("relevance_score", 0.0)
                }
                llm_formatted_data["articles"].append(formatted_article)
            
            json_content = json.dumps(llm_formatted_data, indent=2, ensure_ascii=False)
            
            # 2. Save JSON format
            json_key = f"{date_slug}/{date_slug}_combined.json"
            self.s3_client.put_object(
                Bucket=self.combined_bucket,
                Key=json_key,
                Body=json_content.encode('utf-8'),
                ContentType='application/json'
            )
            paths['json'] = f"s3://{self.combined_bucket}/{json_key}"
            logger.info(f"JSON saved to {paths['json']}")
            
            # 3. Save Markdown format (LLM-friendly)
            markdown_content = self.create_markdown_summary(combined_data)
            md_key = f"{date_slug}/{date_slug}_report.md"
            self.s3_client.put_object(
                Bucket=self.combined_bucket,
                Key=md_key,
                Body=markdown_content.encode('utf-8'),
                ContentType='text/markdown'
            )
            paths['markdown'] = f"s3://{self.combined_bucket}/{md_key}"
            logger.info(f"Markdown saved to {paths['markdown']}")
            
            # 4. Save CSV format (for spreadsheet analysis)
            csv_content = self.create_csv_export(combined_data)
            csv_key = f"{date_slug}/{date_slug}_data.csv"
            self.s3_client.put_object(
                Bucket=self.combined_bucket,
                Key=csv_key,
                Body=csv_content.encode('utf-8'),
                ContentType='text/csv'
            )
            paths['csv'] = f"s3://{self.combined_bucket}/{csv_key}"
            logger.info(f"CSV saved to {paths['csv']}")
            
            # 5. Create ZIP archive with all formats
            if create_zip:
                zip_buffer = BytesIO()
                with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
                    zip_file.writestr(f"{date_slug}_combined.json", json_content)
                    zip_file.writestr(f"{date_slug}_report.md", markdown_content)
                    zip_file.writestr(f"{date_slug}_data.csv", csv_content)
                    
                    # Add a README
                    readme = f"""# Ransomware Intelligence Data - {date_str}

This archive contains ransomware intelligence data in multiple formats:

1. **{date_slug}_combined.json** - Structured JSON data for programmatic analysis
2. **{date_slug}_report.md** - Markdown report for LLM analysis  
3. **{date_slug}_data.csv** - CSV export for spreadsheet analysis

## Summary
- Total Articles: {combined_data['summary']['total_articles']}
- Websites: {len(combined_data['summary']['websites'])}
- Average Relevance Score: {combined_data['summary']['average_relevance_score']:.2f}

Generated: {combined_data['created_at']}

## How to Use

### For LLM Analysis (ChatGPT, Claude, etc.)
Open `{date_slug}_report.md` and paste its contents into your LLM of choice.

### For Spreadsheet Analysis
Open `{date_slug}_data.csv` in Excel, Google Sheets, or any spreadsheet application.

### For Programmatic Analysis
Parse `{date_slug}_combined.json` with your code.
"""
                    zip_file.writestr('README.md', readme)
                
                zip_key = f"{date_slug}/{date_slug}_complete.zip"
                self.s3_client.put_object(
                    Bucket=self.combined_bucket,
                    Key=zip_key,
                    Body=zip_buffer.getvalue(),
                    ContentType='application/zip'
                )
                paths['zip'] = f"s3://{self.combined_bucket}/{zip_key}"
                logger.info(f"ZIP archive saved to {paths['zip']}")
            
            return paths
            
        except ClientError as e:
            logger.error(f"Error saving combined data: {e}")
            raise


def main():
    """Main function"""
    try:
        # Get configuration from environment or command line
        source_bucket = os.getenv("S3_BUCKET_PROCESSED", "ransom-analysis-processed")
        combined_bucket = os.getenv("S3_BUCKET_COMBINED", "ransom-analysis-combined")
        region = os.getenv("AWS_REGION", "us-east-1")
        create_zip = os.getenv("CREATE_ZIP", "true").lower() == "true"
        
        # Get date from command line argument or use yesterday
        if len(sys.argv) > 1:
            date_str = sys.argv[1]
        else:
            # Use yesterday's date
            yesterday = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0)
            yesterday = yesterday.replace(day=yesterday.day - 1)
            date_str = yesterday.strftime("%Y/%m/%d")
        
        logger.info(f"Combining data for date: {date_str}")
        logger.info(f"Source bucket: {source_bucket}")
        logger.info(f"Combined bucket: {combined_bucket}")
        logger.info(f"Region: {region}")
        logger.info(f"Create ZIP: {create_zip}")
        
        # Initialize combiner
        combiner = DailyDataCombiner(source_bucket, combined_bucket, region)
        
        # Combine data
        combined_data = combiner.combine_daily_data(date_str)
        
        if combined_data["total_files"] == 0:
            logger.warning(f"No data found for date {date_str}")
            result = {
                "status": "no_data",
                "date": date_str,
                "message": "No files found for the specified date"
            }
        else:
            # Save combined data in multiple formats
            paths = combiner.save_combined_data_multi_format(combined_data, date_str, create_zip)
            
            result = {
                "status": "success",
                "date": date_str,
                "total_files": combined_data["total_files"],
                "total_articles": combined_data["summary"]["total_articles"],
                "websites": combined_data["summary"]["websites"],
                "features_found": combined_data["summary"]["features_found"],
                "average_relevance_score": combined_data["summary"]["average_relevance_score"],
                "output_paths": paths
            }
        
        print(json.dumps(result, indent=2))
        
    except Exception as e:
        logger.error(f"Error combining daily data: {e}")
        result = {
            "status": "error",
            "error": str(e),
            "timestamp": datetime.now(timezone.utc).isoformat()
        }
        print(json.dumps(result, indent=2))
        sys.exit(1)


if __name__ == "__main__":
    main()
