"""
HTML Report Generator.

Generates interactive HTML reports from analysis results.
"""

from typing import Dict, Any, Optional
from pathlib import Path
import base64
from io import BytesIO
import numpy as np
from PIL import Image

from jinja2 import Template


class ReportGenerator:
    """
    Generate HTML reports from optik analysis results.
    
    Usage:
        generator = ReportGenerator()
        generator.generate(results, "report.html")
    """
    
    def __init__(self):
        self._template = self._get_template()
    
    def generate(
        self, 
        results: Dict[str, Any], 
        output_path: str,
        title: str = "optik Analysis Report"
    ) -> None:
        """
        Generate HTML report from analysis results.
        
        Args:
            results: Dictionary with analysis results from Optik.analyze()
            output_path: Path to save HTML file
            title: Report title
        """
        # Process results for template
        template_data = self._prepare_template_data(results, title)
        
        # Render template
        html = self._template.render(**template_data)
        
        # Write output
        Path(output_path).write_text(html, encoding='utf-8')
    
    def _prepare_template_data(
        self, 
        results: Dict[str, Any],
        title: str
    ) -> Dict[str, Any]:
        """Prepare data for template rendering."""
        data = {
            'title': title,
            'image_path': results.get('image_path', ''),
            'class_names': results.get('class_names', []),
            'num_detections': len(results.get('predictions', [])),
            'detections': [],
            'detection_flow_results': [],
            'grid_trace_results': None,
        }
        
        # Process detections
        for det in results.get('predictions', []):
            data['detections'].append({
                'class_name': det.class_name,
                'confidence': f"{det.confidence:.2%}",
                'bbox': [f"{x:.1f}" for x in det.bbox],
            })
        
        # Process detection flow results
        if 'detection_flow' in results:
            flow_results = results['detection_flow']
            if not isinstance(flow_results, list):
                flow_results = [flow_results]
            
            for flow in flow_results:
                if hasattr(flow, 'final_attribution'):
                    data['detection_flow_results'].append({
                        'detection_idx': flow.detection_idx,
                        'class_name': flow.detection.class_name,
                        'heatmap_base64': self._tensor_to_heatmap_base64(
                            flow.final_attribution
                        ),
                        'scale_contributions': flow.scale_contributions,
                    })
        
        # Process grid trace results
        if 'grid_trace' in results:
            trace = results['grid_trace']
            if hasattr(trace, 'fused_activation'):
                data['grid_trace_results'] = {
                    'heatmap_base64': self._tensor_to_heatmap_base64(
                        trace.fused_activation
                    ),
                    'winning_cells': len(trace.winning_cells),
                    'runner_ups': len(trace.runner_up_cells),
                    'detections_per_scale': trace.detections_per_scale,
                }
        
        return data
    
    def _tensor_to_heatmap_base64(self, tensor) -> str:
        """Convert tensor heatmap to base64 encoded PNG."""
        if tensor is None or tensor.numel() <= 1:
            return ""
        
        # Convert to numpy
        if hasattr(tensor, 'cpu'):
            arr = tensor.cpu().numpy()
        else:
            arr = np.array(tensor)
        
        # Normalize to 0-255
        arr = (arr * 255).clip(0, 255).astype(np.uint8)
        
        # Apply colormap (simple red-yellow gradient)
        h, w = arr.shape[:2] if len(arr.shape) >= 2 else (1, arr.shape[0])
        arr = arr.reshape(h, w) if len(arr.shape) >= 2 else arr.reshape(1, -1)
        
        # Create RGB heatmap
        rgb = np.zeros((h, w, 3), dtype=np.uint8)
        rgb[:, :, 0] = arr  # Red channel
        rgb[:, :, 1] = (arr * 0.5).astype(np.uint8)  # Green
        rgb[:, :, 2] = 0  # Blue
        
        # Convert to PIL and then base64
        img = Image.fromarray(rgb)
        buffer = BytesIO()
        img.save(buffer, format='PNG')
        
        return base64.b64encode(buffer.getvalue()).decode('utf-8')
    
    def _get_template(self) -> Template:
        """Get Jinja2 template for HTML report."""
        template_str = '''
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>{{ title }}</title>
    <style>
        :root {
            --bg-primary: #0f172a;
            --bg-secondary: #1e293b;
            --text-primary: #f1f5f9;
            --text-secondary: #94a3b8;
            --accent: #3b82f6;
            --accent-hover: #60a5fa;
            --success: #10b981;
            --warning: #f59e0b;
        }
        
        * {
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }
        
        body {
            font-family: 'Segoe UI', system-ui, sans-serif;
            background: var(--bg-primary);
            color: var(--text-primary);
            line-height: 1.6;
            padding: 2rem;
        }
        
        .container {
            max-width: 1400px;
            margin: 0 auto;
        }
        
        header {
            text-align: center;
            margin-bottom: 3rem;
            padding-bottom: 2rem;
            border-bottom: 1px solid var(--bg-secondary);
        }
        
        h1 {
            font-size: 2.5rem;
            margin-bottom: 0.5rem;
            background: linear-gradient(135deg, var(--accent), var(--success));
            -webkit-background-clip: text;
            -webkit-text-fill-color: transparent;
        }
        
        h2 {
            font-size: 1.5rem;
            margin: 2rem 0 1rem;
            color: var(--accent);
        }
        
        .card {
            background: var(--bg-secondary);
            border-radius: 12px;
            padding: 1.5rem;
            margin-bottom: 1.5rem;
        }
        
        .stats-grid {
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
            gap: 1rem;
            margin-bottom: 2rem;
        }
        
        .stat-card {
            background: var(--bg-secondary);
            border-radius: 8px;
            padding: 1.5rem;
            text-align: center;
        }
        
        .stat-value {
            font-size: 2rem;
            font-weight: bold;
            color: var(--accent);
        }
        
        .stat-label {
            color: var(--text-secondary);
            font-size: 0.9rem;
        }
        
        .detections-table {
            width: 100%;
            border-collapse: collapse;
        }
        
        .detections-table th,
        .detections-table td {
            padding: 0.75rem 1rem;
            text-align: left;
            border-bottom: 1px solid var(--bg-primary);
        }
        
        .detections-table th {
            background: var(--bg-primary);
            color: var(--text-secondary);
            font-weight: 500;
        }
        
        .heatmap-container {
            display: flex;
            flex-wrap: wrap;
            gap: 1rem;
        }
        
        .heatmap-card {
            flex: 1;
            min-width: 300px;
            background: var(--bg-secondary);
            border-radius: 8px;
            padding: 1rem;
        }
        
        .heatmap-card img {
            width: 100%;
            height: auto;
            border-radius: 4px;
        }
        
        .scale-bar {
            display: flex;
            gap: 0.5rem;
            margin-top: 0.5rem;
        }
        
        .scale-item {
            flex: 1;
            padding: 0.5rem;
            background: var(--bg-primary);
            border-radius: 4px;
            text-align: center;
            font-size: 0.85rem;
        }
        
        .scale-item span {
            display: block;
            font-weight: bold;
            color: var(--accent);
        }
        
        footer {
            text-align: center;
            margin-top: 3rem;
            padding-top: 2rem;
            border-top: 1px solid var(--bg-secondary);
            color: var(--text-secondary);
        }
    </style>
</head>
<body>
    <div class="container">
        <header>
            <h1>🔬 {{ title }}</h1>
            <p style="color: var(--text-secondary);">Generated by optik - AI Model Interpretability Framework</p>
        </header>
        
        <section>
            <div class="stats-grid">
                <div class="stat-card">
                    <div class="stat-value">{{ num_detections }}</div>
                    <div class="stat-label">Detections</div>
                </div>
                <div class="stat-card">
                    <div class="stat-value">{{ class_names|length }}</div>
                    <div class="stat-label">Classes</div>
                </div>
                {% if grid_trace_results %}
                <div class="stat-card">
                    <div class="stat-value">{{ grid_trace_results.winning_cells }}</div>
                    <div class="stat-label">Active Cells</div>
                </div>
                <div class="stat-card">
                    <div class="stat-value">{{ grid_trace_results.runner_ups }}</div>
                    <div class="stat-label">Runner-ups (NMS)</div>
                </div>
                {% endif %}
            </div>
        </section>
        
        <section>
            <h2>📋 Detections</h2>
            <div class="card">
                <table class="detections-table">
                    <thead>
                        <tr>
                            <th>#</th>
                            <th>Class</th>
                            <th>Confidence</th>
                            <th>Bounding Box</th>
                        </tr>
                    </thead>
                    <tbody>
                        {% for det in detections %}
                        <tr>
                            <td>{{ loop.index }}</td>
                            <td>{{ det.class_name }}</td>
                            <td>{{ det.confidence }}</td>
                            <td>[{{ det.bbox|join(', ') }}]</td>
                        </tr>
                        {% endfor %}
                    </tbody>
                </table>
            </div>
        </section>
        
        {% if detection_flow_results %}
        <section>
            <h2>🔥 Detection Flow Attribution</h2>
            <p style="color: var(--text-secondary); margin-bottom: 1rem;">
                What image regions contributed to each detection
            </p>
            <div class="heatmap-container">
                {% for flow in detection_flow_results %}
                <div class="heatmap-card">
                    <h3>Detection {{ flow.detection_idx + 1 }}: {{ flow.class_name }}</h3>
                    {% if flow.heatmap_base64 %}
                    <img src="data:image/png;base64,{{ flow.heatmap_base64 }}" alt="Heatmap">
                    {% endif %}
                    <div class="scale-bar">
                        {% for scale, contrib in flow.scale_contributions.items() %}
                        <div class="scale-item">
                            {{ scale|upper }}
                            <span>{{ "%.0f"|format(contrib * 100) }}%</span>
                        </div>
                        {% endfor %}
                    </div>
                </div>
                {% endfor %}
            </div>
        </section>
        {% endif %}
        
        {% if grid_trace_results %}
        <section>
            <h2>📊 Grid Cell Activation</h2>
            <p style="color: var(--text-secondary); margin-bottom: 1rem;">
                Which FPN grid cells activated for detections
            </p>
            <div class="card">
                {% if grid_trace_results.heatmap_base64 %}
                <img src="data:image/png;base64,{{ grid_trace_results.heatmap_base64 }}" 
                     alt="Grid Activation" style="max-width: 500px;">
                {% endif %}
                <div class="scale-bar" style="max-width: 500px; margin-top: 1rem;">
                    {% for scale, count in grid_trace_results.detections_per_scale.items() %}
                    <div class="scale-item">
                        {{ scale|upper }}
                        <span>{{ count }} dets</span>
                    </div>
                    {% endfor %}
                </div>
            </div>
        </section>
        {% endif %}
        
        <footer>
            <p>Generated by <strong>optik</strong> v0.1.0</p>
            <p>The optical lens for AI model interpretability</p>
        </footer>
    </div>
</body>
</html>
'''
        return Template(template_str)
