"""
ZIP Bundler - Package all cleaning outputs into a single ZIP file.
"""

import zipfile
from io import BytesIO
from pathlib import Path
from typing import Dict, Optional
from datetime import datetime
import pandas as pd


class ZipBundler:
    """Bundles cleaned data, reports, and schemas into a downloadable ZIP file."""
    
    def __init__(self):
        self.buffer = BytesIO()
        
    def create_bundle(
        self,
        cleaned_files: Dict[str, pd.DataFrame],
        merged_df: Optional[pd.DataFrame],
        html_report: str,
        schema_json: Optional[str] = None
    ) -> BytesIO:
        """
        Create a ZIP bundle with all outputs.
        
        Args:
            cleaned_files: Dictionary of filename -> cleaned DataFrame
            merged_df: Merged dataset (if available)
            html_report: HTML report content
            schema_json: JSON schema content (optional)
            
        Returns:
            BytesIO buffer containing the ZIP file
        """
        with zipfile.ZipFile(self.buffer, 'w', zipfile.ZIP_DEFLATED) as zf:
            # Add cleaned files
            for filename, df in cleaned_files.items():
                csv_buffer = BytesIO()
                df.to_csv(csv_buffer, index=False)
                csv_buffer.seek(0)
                zf.writestr(f"cleaned_files/{filename}", csv_buffer.getvalue())
            
            # Add merged dataset
            if merged_df is not None:
                merged_buffer = BytesIO()
                merged_df.to_csv(merged_buffer, index=False)
                merged_buffer.seek(0)
                zf.writestr("merged_dataset.csv", merged_buffer.getvalue())
            
            # Add HTML report
            zf.writestr("cleaning_report.html", html_report)
            
            # Add schema JSON
            if schema_json:
                zf.writestr("schema.json", schema_json)
            
            # Add README
            readme = self._generate_readme(
                num_files=len(cleaned_files),
                has_merged=merged_df is not None,
                has_schema=schema_json is not None
            )
            zf.writestr("README.txt", readme)
        
        self.buffer.seek(0)
        return self.buffer
    
    def _generate_readme(
        self,
        num_files: int,
        has_merged: bool,
        has_schema: bool
    ) -> str:
        """Generate a README for the ZIP bundle."""
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        
        readme = f"""
DataShaper AI - Cleaning Results
================================

Generated: {timestamp}

Contents:
---------
"""
        
        readme += f"✓ {num_files} cleaned CSV file(s) in 'cleaned_files/' folder\n"
        
        if has_merged:
            readme += "✓ merged_dataset.csv - All files merged into one dataset\n"
        
        readme += "✓ cleaning_report.html - Detailed cleaning report (open in browser)\n"
        
        if has_schema:
            readme += "✓ schema.json - Machine-readable data schema\n"
        
        readme += """
How to Use:
-----------
1. Open cleaning_report.html to see what was cleaned
2. Use individual files in cleaned_files/ folder
3. Or use merged_dataset.csv for combined data
"""
        
        if has_schema:
            readme += "4. Use schema.json for ML pipelines or data validation\n"
        
        readme += """
DataShaper AI - Deterministic Data Cleaning
No AI hallucinations, just trustworthy results.
"""
        
        return readme
    
    @staticmethod
    def get_bundle_filename() -> str:
        """Generate a timestamped filename for the ZIP bundle."""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        return f"datashaper_cleaned_{timestamp}.zip"
