#!/usr/bin/env python3
"""
Storage Saturation Report Export Script
Converts JSON report data to PDF or JPG format using weasyprint and Pillow
"""

import sys
import json
import argparse
import os
from datetime import datetime

try:
    from weasyprint import HTML, CSS
except ImportError:
    print("Error: weasyprint module not installed. Please install with: pip3 install weasyprint", file=sys.stderr)
    sys.exit(1)

try:
    from PIL import Image
    from pdf2image import convert_from_path
except ImportError:
    print("Error: Pillow or pdf2image module not installed. Please install with: pip3 install Pillow pdf2image", file=sys.stderr)
    sys.exit(1)


def generate_html_template(data):
    """Generate HTML report from JSON data"""
    
    title = data.get('title', 'Storage Saturation Report')
    date = data.get('date', datetime.now().strftime('%m/%d/%y, %I:%M %p'))
    services = data.get('services', [])
    
    # Determine color class for available space
    def get_available_class(available_bytes):
        if available_bytes == 0:
            return 'text-muted'
        available_gb = available_bytes / (1024 * 1024 * 1024)
        if available_gb < 5:
            return 'text-danger'
        elif available_gb < 50:
            return 'text-warning'
        return 'text-success'
    
    # Determine progress bar class
    def get_progress_class(percent):
        if percent >= 95:
            return 'progress-bar-danger'
        elif percent >= 85:
            return 'progress-bar-warning'
        elif percent >= 70:
            return 'progress-bar-info'
        return 'progress-bar-success'
    
    # Build table rows
    table_rows = ''
    if isinstance(services, dict) and 'error' in services:
        table_rows = f'''
            <tr>
                <td colspan="5" style="text-align: center; padding: 20px; color: #d9534f; font-size: 14px;">
                    <strong>Error:</strong> {services['error']}
                </td>
            </tr>
        '''
    elif isinstance(services, list) and len(services) > 0:
        for service in services:
            host_name = service.get('host_name', 'N/A')
            caption = service.get('service_description', service.get('name', 'N/A'))
            disk_used = service.get('disk_used_display', 'N/A')
            disk_available = service.get('disk_available_display', 'N/A')
            disk_percent = service.get('disk_usage_percent', 0)
            disk_percent_display = service.get('disk_percent_display', 'N/A')
            
            # Get color classes
            available_bytes = service.get('disk_available_bytes', 0)
            available_class = get_available_class(available_bytes)
            progress_class = get_progress_class(disk_percent)
            
            # Build progress bar HTML
            if disk_percent > 0 and disk_percent_display != 'N/A':
                progress_bar = f'''
                    <div style="width: 100%; background-color: #f0f0f0; border-radius: 4px; overflow: hidden;">
                        <div style="width: {disk_percent}%; background-color: {'#d9534f' if disk_percent >= 95 else '#f0ad4e' if disk_percent >= 85 else '#5bc0de' if disk_percent >= 70 else '#5cb85c'}; 
                             height: 20px; text-align: center; line-height: 20px; color: white; font-size: 11px; font-weight: bold;">
                            {disk_percent_display}
                        </div>
                    </div>
                '''
            else:
                progress_bar = '<span>N/A</span>'
            
            # Determine text color for available
            available_color = '#d9534f' if available_class == 'text-danger' else '#f0ad4e' if available_class == 'text-warning' else '#5cb85c'
            
            table_rows += f'''
                <tr>
                    <td>{host_name}</td>
                    <td>{caption}</td>
                    <td>{disk_used}</td>
                    <td style="color: {available_color};">{disk_available}</td>
                    <td>{progress_bar}</td>
                </tr>
            '''
    else:
        table_rows = '''
            <tr>
                <td colspan="5" style="text-align: center; padding: 20px; color: #999;">
                    No data available
                </td>
            </tr>
        '''
    
    html_template = f'''<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <title>{title}</title>
    <style>
        @page {{
            size: letter;
            margin: 0.75in;
        }}
        body {{
            font-family: Arial, sans-serif;
            margin: 0;
            padding: 20px;
            background: #fff;
            color: #333;
        }}
        .report-header {{
            margin-bottom: 30px;
            border-bottom: 2px solid #2c3e50;
            padding-bottom: 15px;
        }}
        .report-header h1 {{
            margin: 0 0 10px 0;
            font-size: 24px;
            font-weight: bold;
            color: #2c3e50;
        }}
        .report-header .subtitle {{
            font-size: 14px;
            color: #666;
            margin: 5px 0;
        }}
        .report-header .report-date {{
            font-size: 12px;
            color: #999;
            margin: 5px 0;
        }}
        .report-header .report-order {{
            font-size: 12px;
            color: #999;
            margin: 5px 0;
        }}
        table {{
            width: 100%;
            border-collapse: collapse;
            margin-top: 20px;
        }}
        th {{
            background-color: #2c3e50;
            color: #ecf0f1;
            font-weight: bold;
            text-align: left;
            padding: 12px 15px;
            border: 1px solid #34495e;
            text-transform: uppercase;
            font-size: 12px;
            letter-spacing: 0.5px;
        }}
        td {{
            padding: 10px 15px;
            border: 1px solid #e0e0e0;
            font-size: 13px;
        }}
        tr:nth-child(even) {{
            background-color: #f9f9f9;
        }}
        tr:last-child td {{
            border-bottom: none;
        }}
        .text-danger {{
            color: #d9534f;
        }}
        .text-warning {{
            color: #f0ad4e;
        }}
        .text-success {{
            color: #5cb85c;
        }}
        .text-muted {{
            color: #999;
        }}
    </style>
</head>
<body>
    <div class="report-header">
        <h1>{title}</h1>
        <p class="subtitle">Summary of Tracked Storage: Volumes</p>
        <p class="report-date">Report Date: {date}</p>
        <p class="report-order">Ordered by: Disk Space Available - Ascending</p>
    </div>
    
    <table>
        <thead>
            <tr>
                <th>DISPLAY NAME</th>
                <th>CAPTION</th>
                <th>DISK SPACE USED</th>
                <th>DISK SPACE AVAILABLE</th>
                <th>PERCENT USED</th>
            </tr>
        </thead>
        <tbody>
            {table_rows}
        </tbody>
    </table>
</body>
</html>
'''
    
    return html_template


def convert_to_pdf(html_content, output_file):
    """Convert HTML to PDF using weasyprint"""
    try:
        HTML(string=html_content).write_pdf(output_file)
        return True
    except Exception as e:
        print(f"Error converting to PDF: {e}", file=sys.stderr)
        return False


def convert_pdf_to_jpg(pdf_file, jpg_file):
    """Convert PDF to JPG using pdf2image"""
    try:
        # Convert PDF to images (first page only)
        images = convert_from_path(pdf_file, first_page=1, last_page=1, dpi=150)
        if images:
            # Save as JPG
            images[0].save(jpg_file, 'JPEG', quality=95)
            return True
        return False
    except Exception as e:
        print(f"Error converting PDF to JPG: {e}", file=sys.stderr)
        return False


def main():
    parser = argparse.ArgumentParser(description='Export Storage Saturation Report to PDF or JPG')
    parser.add_argument('--input', required=True, help='Input JSON file path')
    parser.add_argument('--output', required=True, help='Output file path')
    parser.add_argument('--format', required=True, choices=['pdf', 'jpg'], help='Output format (pdf or jpg)')
    
    args = parser.parse_args()
    
    # Read JSON input
    try:
        with open(args.input, 'r', encoding='utf-8') as f:
            data = json.load(f)
    except FileNotFoundError:
        print(f"Error: Input file not found: {args.input}", file=sys.stderr)
        sys.exit(1)
    except json.JSONDecodeError as e:
        print(f"Error: Invalid JSON in input file: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Error reading input file: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Generate HTML
    html_content = generate_html_template(data)
    
    # Convert to PDF
    if args.format == 'pdf':
        if not convert_to_pdf(html_content, args.output):
            sys.exit(1)
    elif args.format == 'jpg':
        # First convert to PDF, then to JPG
        temp_pdf = args.output.replace('.jpg', '.pdf').replace('.jpeg', '.pdf')
        if not convert_to_pdf(html_content, temp_pdf):
            sys.exit(1)
        
        if not convert_pdf_to_jpg(temp_pdf, args.output):
            # Clean up temp PDF
            if os.path.exists(temp_pdf):
                os.remove(temp_pdf)
            sys.exit(1)
        
        # Clean up temp PDF
        if os.path.exists(temp_pdf):
            os.remove(temp_pdf)
    
    print(f"Successfully created {args.format.upper()}: {args.output}")
    sys.exit(0)


if __name__ == '__main__':
    main()