# Copyright 2025 Thulium Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Benchmark result reporting and export utilities.

This module provides tools for generating formatted reports
from benchmark results in multiple output formats:
- Markdown tables for documentation
- CSV for spreadsheet analysis
- JSON for programmatic consumption
- HTML for web display

The reporting utilities support per-language breakdowns,
model comparisons, and illustrative performance summaries.
"""

from __future__ import annotations

import csv
import io
import json
from datetime import datetime
from pathlib import Path
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Union

from thulium.evaluation.benchmarking import BenchmarkResult
from thulium.evaluation.benchmarking import LanguageResult


def generate_markdown_report(
    result: BenchmarkResult,
    include_samples: bool = False,
    max_samples: int = 20
) -> str:
    """
    Generate a Markdown-formatted benchmark report.
    
    Args:
        result: BenchmarkResult to format.
        include_samples: If True, include per-sample details.
        max_samples: Maximum samples to include in detail section.
        
    Returns:
        Markdown-formatted report string.
    """
    lines = [
        f"# Benchmark Report: {result.config.name}",
        "",
        f"**Generated**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        "",
        f"**Description**: {result.config.description}",
        "",
        "---",
        "",
        "## Summary",
        "",
        "| Metric | Value |",
        "|:-------|------:|",
        f"| Total Samples | {result.num_samples} |",
        f"| Character Error Rate (CER) | {result.aggregate_cer * 100:.2f}% |",
        f"| Word Error Rate (WER) | {result.aggregate_wer * 100:.2f}% |",
        f"| Sequence Error Rate (SER) | {result.aggregate_ser * 100:.2f}% |",
        f"| Average Latency | {result.avg_latency_ms:.2f} ms |",
        f"| Throughput | {result.throughput_samples_per_sec:.2f} samples/sec |",
        f"| Total Time | {result.total_time_s:.2f} s |",
        "",
    ]
    
    # Per-language breakdown
    if result.per_language:
        lines.extend([
            "## Results by Language",
            "",
            "| Language | Samples | CER (%) | WER (%) | SER (%) | Latency (ms) |",
            "|:---------|--------:|--------:|--------:|--------:|-------------:|",
        ])
        
        for lang_result in result.per_language:
            lines.append(
                f"| {lang_result.language} | {lang_result.num_samples} | "
                f"{lang_result.cer * 100:.2f} | {lang_result.wer * 100:.2f} | "
                f"{lang_result.ser * 100:.2f} | {lang_result.avg_latency_ms:.2f} |"
            )
        
        lines.append("")
    
    # Configuration details
    lines.extend([
        "## Configuration",
        "",
        f"- **Model Config**: `{result.config.model_config}`",
        f"- **Device**: {result.config.device}",
        f"- **Batch Size**: {result.config.batch_size}",
        f"- **Seed**: {result.config.seed}",
        "",
    ])
    
    # Decoding parameters
    if result.config.decoding:
        lines.extend([
            "### Decoding Parameters",
            "",
        ])
        for key, value in result.config.decoding.items():
            lines.append(f"- **{key}**: {value}")
        lines.append("")
    
    # Sample details
    if include_samples and result.sample_results:
        lines.extend([
            "## Sample Results",
            "",
            "| ID | Language | CER | WER | Reference | Hypothesis |",
            "|:---|:---------|----:|----:|:----------|:-----------|",
        ])
        
        for sample in result.sample_results[:max_samples]:
            ref_truncated = sample.reference[:30] + "..." if len(sample.reference) > 30 else sample.reference
            hyp_truncated = sample.hypothesis[:30] + "..." if len(sample.hypothesis) > 30 else sample.hypothesis
            lines.append(
                f"| {sample.sample_id} | {sample.language or '-'} | "
                f"{sample.cer:.3f} | {sample.wer:.3f} | "
                f"{ref_truncated} | {hyp_truncated} |"
            )
        
        if len(result.sample_results) > max_samples:
            lines.append(f"| ... | ... | ... | ... | *({len(result.sample_results) - max_samples} more samples)* | ... |")
        
        lines.append("")
    
    # Footer
    lines.extend([
        "---",
        "",
        "*Report generated by Thulium Evaluation Suite*",
    ])
    
    return '\n'.join(lines)


def generate_csv_report(
    result: BenchmarkResult,
    include_samples: bool = True
) -> str:
    """
    Generate a CSV-formatted benchmark report.
    
    Args:
        result: BenchmarkResult to format.
        include_samples: If True, include per-sample rows.
        
    Returns:
        CSV-formatted report string.
    """
    output = io.StringIO()
    
    if include_samples and result.sample_results:
        # Sample-level CSV
        writer = csv.writer(output)
        writer.writerow([
            'sample_id', 'language', 'cer', 'wer', 'latency_ms',
            'reference', 'hypothesis'
        ])
        
        for sample in result.sample_results:
            writer.writerow([
                sample.sample_id,
                sample.language or '',
                f"{sample.cer:.6f}",
                f"{sample.wer:.6f}",
                f"{sample.latency_ms:.3f}",
                sample.reference,
                sample.hypothesis
            ])
    else:
        # Aggregate-level CSV
        writer = csv.writer(output)
        writer.writerow([
            'config_name', 'num_samples', 'cer', 'wer', 'ser',
            'avg_latency_ms', 'throughput'
        ])
        writer.writerow([
            result.config.name,
            result.num_samples,
            f"{result.aggregate_cer:.6f}",
            f"{result.aggregate_wer:.6f}",
            f"{result.aggregate_ser:.6f}",
            f"{result.avg_latency_ms:.3f}",
            f"{result.throughput_samples_per_sec:.3f}"
        ])
    
    return output.getvalue()


def generate_json_report(
    result: BenchmarkResult,
    include_samples: bool = True,
    indent: int = 2
) -> str:
    """
    Generate a JSON-formatted benchmark report.
    
    Args:
        result: BenchmarkResult to format.
        include_samples: If True, include per-sample details.
        indent: JSON indentation level.
        
    Returns:
        JSON-formatted report string.
    """
    data = result.to_dict()
    if not include_samples:
        data.pop('sample_results', None)
    
    data['generated_at'] = datetime.now().isoformat()
    
    return json.dumps(data, indent=indent, ensure_ascii=False)


def generate_comparison_table(
    results: List[BenchmarkResult],
    format: str = 'markdown'
) -> str:
    """
    Generate a comparison table for multiple benchmark results.
    
    Args:
        results: List of BenchmarkResult objects.
        format: Output format ('markdown', 'csv', 'html').
        
    Returns:
        Formatted comparison table.
    """
    if format == 'markdown':
        lines = [
            "# Model Comparison",
            "",
            "> Note: These results are based on development benchmarks. "
            "Actual performance may vary depending on dataset and configuration.",
            "",
            "| Model | CER (%) | WER (%) | Latency (ms) | Throughput (samples/s) |",
            "|:------|--------:|--------:|-------------:|-----------------------:|"
        ]
        
        for r in results:
            lines.append(
                f"| {r.config.name} | {r.aggregate_cer * 100:.2f} | "
                f"{r.aggregate_wer * 100:.2f} | {r.avg_latency_ms:.1f} | "
                f"{r.throughput_samples_per_sec:.1f} |"
            )
        
        return '\n'.join(lines)
    
    elif format == 'csv':
        output = io.StringIO()
        writer = csv.writer(output)
        writer.writerow(['model', 'cer', 'wer', 'latency_ms', 'throughput'])
        for r in results:
            writer.writerow([
                r.config.name,
                f"{r.aggregate_cer:.6f}",
                f"{r.aggregate_wer:.6f}",
                f"{r.avg_latency_ms:.3f}",
                f"{r.throughput_samples_per_sec:.3f}"
            ])
        return output.getvalue()
    
    elif format == 'html':
        rows = []
        for r in results:
            rows.append(
                f"<tr><td>{r.config.name}</td>"
                f"<td>{r.aggregate_cer * 100:.2f}%</td>"
                f"<td>{r.aggregate_wer * 100:.2f}%</td>"
                f"<td>{r.avg_latency_ms:.1f}</td>"
                f"<td>{r.throughput_samples_per_sec:.1f}</td></tr>"
            )
        
        return f"""<table>
<thead>
<tr><th>Model</th><th>CER</th><th>WER</th><th>Latency (ms)</th><th>Throughput</th></tr>
</thead>
<tbody>
{''.join(rows)}
</tbody>
</table>"""
    
    else:
        raise ValueError(f"Unknown format: {format}")


def save_report(
    result: BenchmarkResult,
    output_path: Union[str, Path],
    format: str = 'auto',
    include_samples: bool = True
) -> None:
    """
    Save benchmark report to file.
    
    Args:
        result: BenchmarkResult to save.
        output_path: Output file path.
        format: Output format. 'auto' infers from extension.
        include_samples: Whether to include per-sample details.
    """
    output_path = Path(output_path)
    
    if format == 'auto':
        ext = output_path.suffix.lower()
        format_map = {'.md': 'markdown', '.csv': 'csv', '.json': 'json', '.html': 'html'}
        format = format_map.get(ext, 'json')
    
    if format == 'markdown':
        content = generate_markdown_report(result, include_samples)
    elif format == 'csv':
        content = generate_csv_report(result, include_samples)
    elif format == 'json':
        content = generate_json_report(result, include_samples)
    else:
        raise ValueError(f"Unknown format: {format}")
    
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(content)
