# smoltrace/cards.py
"""Dataset card generation for SMOLTRACE datasets with branding and documentation."""

from datetime import datetime
from typing import Optional

# SMOLTRACE branding constants
SMOLTRACE_LOGO_URL = (
    "https://raw.githubusercontent.com/Mandark-droid/SMOLTRACE/main/.github/images/Logo.png"
)
SMOLTRACE_REPO_URL = "https://github.com/Mandark-droid/SMOLTRACE"
SMOLTRACE_PYPI_URL = "https://pypi.org/project/smoltrace/"
SMOLTRACE_DOCS_URL = "https://github.com/Mandark-droid/SMOLTRACE#readme"


def _get_header() -> str:
    """Returns the common header with SMOLTRACE branding."""
    return f"""---
license: agpl-3.0
tags:
  - smoltrace
  - smolagents
  - evaluation
  - benchmark
  - llm
  - agents
---

<div align="center">
  <img src="{SMOLTRACE_LOGO_URL}" alt="SMOLTRACE Logo" width="400"/>

  <h3>Tiny Agents. Total Visibility.</h3>

  <p>
    <a href="{SMOLTRACE_REPO_URL}"><img src="https://img.shields.io/badge/GitHub-SMOLTRACE-blue?logo=github" alt="GitHub"></a>
    <a href="{SMOLTRACE_PYPI_URL}"><img src="https://img.shields.io/pypi/v/smoltrace?color=green" alt="PyPI"></a>
    <a href="{SMOLTRACE_DOCS_URL}"><img src="https://img.shields.io/badge/docs-readme-orange" alt="Documentation"></a>
  </p>
</div>

---

"""


def _get_footer() -> str:
    """Returns the common footer with links and attribution."""
    return f"""
---

## About SMOLTRACE

**SMOLTRACE** is a comprehensive benchmarking and evaluation framework for [Smolagents](https://huggingface.co/docs/smolagents) - HuggingFace's lightweight agent library.

### Key Features
- Automated agent evaluation with customizable test cases
- OpenTelemetry-based tracing for detailed execution insights
- GPU metrics collection (utilization, memory, temperature, power)
- CO2 emissions and power cost tracking
- Leaderboard aggregation and comparison

### Quick Links
- [GitHub Repository]({SMOLTRACE_REPO_URL})
- [PyPI Package]({SMOLTRACE_PYPI_URL})
- [Documentation]({SMOLTRACE_DOCS_URL})
- [Report Issues]({SMOLTRACE_REPO_URL}/issues)

### Installation

```bash
pip install smoltrace
```

### Citation

If you use SMOLTRACE in your research, please cite:

```bibtex
@software{{smoltrace,
  title = {{SMOLTRACE: Benchmarking Framework for Smolagents}},
  author = {{Thakkar, Kshitij}},
  url = {{{SMOLTRACE_REPO_URL}}},
  year = {{2025}}
}}
```

---

<div align="center">
  <sub>Generated by <a href="{SMOLTRACE_REPO_URL}">SMOLTRACE</a></sub>
</div>
"""


def generate_results_card(
    model_name: str,
    run_id: str,
    num_results: int,
    agent_type: str = "both",
    dataset_used: Optional[str] = None,
) -> str:
    """
    Generate a dataset card for evaluation results.

    Args:
        model_name: Name of the evaluated model
        run_id: Unique run identifier
        num_results: Number of test results
        agent_type: Agent type used ("tool", "code", or "both")
        dataset_used: Source dataset used for evaluation

    Returns:
        Markdown content for the dataset card
    """
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC")

    return f"""{_get_header()}
# SMOLTRACE Evaluation Results

This dataset contains evaluation results from a SMOLTRACE benchmark run.

## Dataset Information

| Field | Value |
|-------|-------|
| **Model** | `{model_name}` |
| **Run ID** | `{run_id}` |
| **Agent Type** | `{agent_type}` |
| **Total Tests** | {num_results} |
| **Generated** | {timestamp} |
| **Source Dataset** | {dataset_used or "N/A"} |

## Schema

| Column | Type | Description |
|--------|------|-------------|
| `model` | string | Model identifier |
| `evaluation_date` | string | ISO timestamp of evaluation |
| `task_id` | string | Unique test case identifier |
| `agent_type` | string | "tool" or "code" agent type |
| `difficulty` | string | Test difficulty level |
| `prompt` | string | Test prompt/question |
| `success` | bool | Whether the test passed |
| `tool_called` | bool | Whether a tool was invoked |
| `correct_tool` | bool | Whether the correct tool was used |
| `final_answer_called` | bool | Whether final_answer was called |
| `response_correct` | bool | Whether the response was correct |
| `tools_used` | string | Comma-separated list of tools used |
| `steps` | int | Number of agent steps taken |
| `response` | string | Agent's final response |
| `error` | string | Error message if failed |
| `trace_id` | string | OpenTelemetry trace ID |
| `execution_time_ms` | float | Execution time in milliseconds |
| `total_tokens` | int | Total tokens consumed |
| `cost_usd` | float | API cost in USD |
| `enhanced_trace_info` | string | JSON with detailed trace data |

## Usage

```python
from datasets import load_dataset

# Load the results dataset
ds = load_dataset("{run_id.split('/')[0] if '/' in run_id else 'YOUR_USERNAME'}/smoltrace-results-TIMESTAMP")

# Filter successful tests
successful = ds.filter(lambda x: x['success'])

# Calculate success rate
success_rate = sum(1 for r in ds['train'] if r['success']) / len(ds['train']) * 100
print(f"Success Rate: {{success_rate:.2f}}%")
```

## Related Datasets

This evaluation run also generated:
- **Traces Dataset**: Detailed OpenTelemetry execution traces
- **Metrics Dataset**: GPU utilization and environmental metrics
- **Leaderboard**: Aggregated metrics for model comparison
{_get_footer()}"""


def generate_traces_card(
    model_name: str,
    run_id: str,
    num_traces: int,
) -> str:
    """
    Generate a dataset card for OpenTelemetry traces.

    Args:
        model_name: Name of the evaluated model
        run_id: Unique run identifier
        num_traces: Number of trace records

    Returns:
        Markdown content for the dataset card
    """
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC")

    return f"""{_get_header()}
# SMOLTRACE Execution Traces

This dataset contains OpenTelemetry execution traces from a SMOLTRACE benchmark run.

## Dataset Information

| Field | Value |
|-------|-------|
| **Model** | `{model_name}` |
| **Run ID** | `{run_id}` |
| **Total Traces** | {num_traces} |
| **Generated** | {timestamp} |
| **Format** | OpenTelemetry-compatible |

## Schema

| Column | Type | Description |
|--------|------|-------------|
| `trace_id` | string | Unique trace identifier |
| `span_id` | string | Span identifier within trace |
| `parent_span_id` | string | Parent span for hierarchy |
| `span_name` | string | Name of the operation |
| `start_time` | string | ISO timestamp of span start |
| `end_time` | string | ISO timestamp of span end |
| `duration_ms` | float | Span duration in milliseconds |
| `status` | string | Span status (OK, ERROR) |
| `attributes` | string | JSON with span attributes |
| `events` | string | JSON with span events |
| `total_tokens` | int | Tokens used in this span |
| `input_tokens` | int | Input/prompt tokens |
| `output_tokens` | int | Output/completion tokens |
| `total_cost_usd` | float | Cost for this span |
| `total_duration_ms` | float | Total duration including children |

## Understanding Traces

SMOLTRACE uses OpenTelemetry to capture detailed execution traces:

```
agent.run (root span)
├── agent.step (step 1)
│   ├── llm.call (model inference)
│   └── tool.execute (tool invocation)
├── agent.step (step 2)
│   └── llm.call
└── agent.finalize
```

## Usage

```python
from datasets import load_dataset
import json

# Load traces
ds = load_dataset("{run_id.split('/')[0] if '/' in run_id else 'YOUR_USERNAME'}/smoltrace-traces-TIMESTAMP")

# Analyze execution patterns
for trace in ds['train']:
    attrs = json.loads(trace['attributes']) if trace['attributes'] else {{}}
    print(f"Trace {{trace['trace_id']}}: {{trace['span_name']}} - {{trace['duration_ms']}}ms")
```

## Related Datasets

This evaluation run also generated:
- **Results Dataset**: Pass/fail outcomes for each test case
- **Metrics Dataset**: GPU utilization and environmental metrics
- **Leaderboard**: Aggregated metrics for model comparison
{_get_footer()}"""


def generate_metrics_card(
    model_name: str,
    run_id: str,
    num_metrics: int,
    has_gpu_metrics: bool = True,
) -> str:
    """
    Generate a dataset card for GPU and environmental metrics.

    Args:
        model_name: Name of the evaluated model
        run_id: Unique run identifier
        num_metrics: Number of metric records
        has_gpu_metrics: Whether GPU metrics are available

    Returns:
        Markdown content for the dataset card
    """
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC")
    gpu_status = "Available" if has_gpu_metrics else "N/A (API Model)"

    return f"""{_get_header()}
# SMOLTRACE GPU & Environmental Metrics

This dataset contains time-series GPU metrics and environmental impact data from a SMOLTRACE benchmark run.

## Dataset Information

| Field | Value |
|-------|-------|
| **Model** | `{model_name}` |
| **Run ID** | `{run_id}` |
| **Total Samples** | {num_metrics} |
| **Generated** | {timestamp} |
| **GPU Metrics** | {gpu_status} |

## Schema

| Column | Type | Description |
|--------|------|-------------|
| `run_id` | string | Unique run identifier |
| `timestamp` | string | ISO timestamp of measurement |
| `timestamp_unix_nano` | string | Unix nanosecond timestamp |
| `service_name` | string | Service identifier |
| `gpu_id` | string | GPU device ID |
| `gpu_name` | string | GPU model name |
| `gpu_utilization_percent` | float | GPU compute utilization (0-100%) |
| `gpu_memory_used_mib` | float | GPU memory used (MiB) |
| `gpu_memory_total_mib` | float | Total GPU memory (MiB) |
| `gpu_temperature_celsius` | float | GPU temperature (°C) |
| `gpu_power_watts` | float | GPU power consumption (W) |
| `co2_emissions_gco2e` | float | Cumulative CO2 emissions (gCO2e) |
| `power_cost_usd` | float | Cumulative power cost (USD) |

## Environmental Impact

SMOLTRACE tracks environmental metrics to help you understand the carbon footprint of your AI workloads:

- **CO2 Emissions**: Calculated based on GPU power consumption and regional carbon intensity
- **Power Cost**: Estimated electricity cost based on configurable rates

## Usage

```python
from datasets import load_dataset
import pandas as pd

# Load metrics
ds = load_dataset("{run_id.split('/')[0] if '/' in run_id else 'YOUR_USERNAME'}/smoltrace-metrics-TIMESTAMP")

# Convert to DataFrame for analysis
df = pd.DataFrame(ds['train'])

# Plot GPU utilization over time
import matplotlib.pyplot as plt
plt.plot(df['timestamp'], df['gpu_utilization_percent'])
plt.xlabel('Time')
plt.ylabel('GPU Utilization (%)')
plt.title('GPU Utilization During Evaluation')
plt.show()

# Get total environmental impact
total_co2 = df['co2_emissions_gco2e'].max()
total_cost = df['power_cost_usd'].max()
print(f"Total CO2: {{total_co2:.4f}} gCO2e")
print(f"Total Cost: ${{total_cost:.6f}}")
```

## Related Datasets

This evaluation run also generated:
- **Results Dataset**: Pass/fail outcomes for each test case
- **Traces Dataset**: Detailed OpenTelemetry execution traces
- **Leaderboard**: Aggregated metrics for model comparison
{_get_footer()}"""


def generate_leaderboard_card(username: str) -> str:
    """
    Generate a dataset card for the leaderboard.

    Args:
        username: HuggingFace username

    Returns:
        Markdown content for the dataset card
    """
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC")

    return f"""{_get_header()}
# SMOLTRACE Leaderboard

This dataset contains aggregated evaluation metrics for comparing model performance across SMOLTRACE benchmark runs.

## Dataset Information

| Field | Value |
|-------|-------|
| **Owner** | `{username}` |
| **Updated** | {timestamp} |
| **Purpose** | Model comparison and ranking |

## Schema

### Identification
| Column | Type | Description |
|--------|------|-------------|
| `run_id` | string | Unique run identifier |
| `model` | string | Model name/identifier |
| `agent_type` | string | Agent type ("tool", "code", "both") |
| `provider` | string | Model provider (litellm, openai, etc.) |
| `timestamp` | string | Evaluation timestamp |
| `submitted_by` | string | HuggingFace username |

### Dataset References
| Column | Type | Description |
|--------|------|-------------|
| `results_dataset` | string | Link to results dataset |
| `traces_dataset` | string | Link to traces dataset |
| `metrics_dataset` | string | Link to metrics dataset |
| `dataset_used` | string | Source benchmark dataset |

### Performance Metrics
| Column | Type | Description |
|--------|------|-------------|
| `total_tests` | int | Number of test cases |
| `successful_tests` | int | Passed tests |
| `failed_tests` | int | Failed tests |
| `success_rate` | float | Success percentage (0-100) |
| `avg_steps` | float | Average agent steps per test |
| `avg_duration_ms` | float | Average execution time (ms) |
| `total_duration_ms` | float | Total evaluation time (ms) |
| `total_tokens` | int | Total tokens consumed |
| `avg_tokens_per_test` | int | Average tokens per test |
| `total_cost_usd` | float | Total API cost (USD) |
| `avg_cost_per_test_usd` | float | Average cost per test (USD) |

### Environmental Impact
| Column | Type | Description |
|--------|------|-------------|
| `co2_emissions_g` | float | Total CO2 emissions (gCO2e) |
| `power_cost_total_usd` | float | Total power cost (USD) |

### GPU Metrics (if available)
| Column | Type | Description |
|--------|------|-------------|
| `gpu_utilization_avg` | float | Average GPU utilization (%) |
| `gpu_utilization_max` | float | Peak GPU utilization (%) |
| `gpu_memory_avg_mib` | float | Average GPU memory (MiB) |
| `gpu_memory_max_mib` | float | Peak GPU memory (MiB) |
| `gpu_temperature_avg` | float | Average GPU temperature (°C) |
| `gpu_temperature_max` | float | Peak GPU temperature (°C) |
| `gpu_power_avg_w` | float | Average GPU power (W) |

## Usage

```python
from datasets import load_dataset
import pandas as pd

# Load leaderboard
ds = load_dataset("{username}/smoltrace-leaderboard")
df = pd.DataFrame(ds['train'])

# Rank by success rate
df_ranked = df.sort_values('success_rate', ascending=False)
print(df_ranked[['model', 'success_rate', 'avg_duration_ms', 'total_cost_usd']])

# Compare models
top_models = df_ranked.head(10)
print("Top 10 Models by Success Rate:")
for i, row in top_models.iterrows():
    print(f"  {{row['model']}}: {{row['success_rate']:.1f}}%")
```

## Contributing Results

Run your own evaluations to add to this leaderboard:

```bash
pip install smoltrace
smoltrace-eval --model your-model --provider litellm
```
{_get_footer()}"""


def generate_benchmark_card(
    username: str,
    num_cases: int,
    source_user: str = "kshitijthakkar",
) -> str:
    """
    Generate a dataset card for the benchmark dataset.

    Args:
        username: HuggingFace username (destination)
        num_cases: Number of test cases
        source_user: Original source username

    Returns:
        Markdown content for the dataset card
    """
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC")

    return f"""{_get_header()}
# SMOLTRACE Benchmark v1

The official comprehensive benchmark dataset for evaluating Smolagents with SMOLTRACE.

## Dataset Information

| Field | Value |
|-------|-------|
| **Owner** | `{username}` |
| **Test Cases** | {num_cases} |
| **Source** | `{source_user}/smoltrace-benchmark-v1` |
| **Copied** | {timestamp} |

## Overview

This benchmark contains {num_cases} carefully curated test cases designed to evaluate:

- **Tool Usage**: Correct tool selection and invocation
- **Code Generation**: Python code writing and execution
- **Reasoning**: Multi-step problem solving
- **Web Search**: Information retrieval capabilities
- **Math**: Numerical calculations
- **General Knowledge**: Factual question answering

## Schema

| Column | Type | Description |
|--------|------|-------------|
| `id` | string | Unique test case identifier |
| `prompt` | string | Test prompt/question |
| `expected_tool` | string | Expected tool to be called |
| `expected_answer` | string | Expected answer (if applicable) |
| `difficulty` | string | easy, medium, or hard |
| `category` | string | Test category |
| `agent_type` | string | Suitable agent type |

## Difficulty Distribution

| Difficulty | Count | Description |
|------------|-------|-------------|
| **Easy** | ~40% | Basic tool calls, simple questions |
| **Medium** | ~40% | Multi-step reasoning, complex queries |
| **Hard** | ~20% | Advanced problems, edge cases |

## Usage

```bash
# Run evaluation with this benchmark
smoltrace-eval --model gpt-4 --dataset-name {username}/smoltrace-benchmark-v1

# Run specific agent type
smoltrace-eval --model claude-3-opus --agent-type tool --dataset-name {username}/smoltrace-benchmark-v1
```

```python
from datasets import load_dataset

# Load benchmark
ds = load_dataset("{username}/smoltrace-benchmark-v1")

# Explore test cases
for case in ds['train']:
    print(f"[{{case['difficulty']}}] {{case['prompt'][:50]}}...")
```
{_get_footer()}"""


def generate_tasks_card(
    username: str,
    num_cases: int,
    source_user: str = "kshitijthakkar",
) -> str:
    """
    Generate a dataset card for the tasks dataset.

    Args:
        username: HuggingFace username (destination)
        num_cases: Number of test cases
        source_user: Original source username

    Returns:
        Markdown content for the dataset card
    """
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC")

    return f"""{_get_header()}
# SMOLTRACE Tasks

A lightweight default task set for quick SMOLTRACE evaluations.

## Dataset Information

| Field | Value |
|-------|-------|
| **Owner** | `{username}` |
| **Test Cases** | {num_cases} |
| **Source** | `{source_user}/smoltrace-tasks` |
| **Copied** | {timestamp} |

## Overview

This is a curated subset of {num_cases} representative test cases, ideal for:

- Quick model validation
- CI/CD pipeline testing
- Development and debugging
- Initial model screening

For comprehensive evaluation, use `smoltrace-benchmark-v1` instead.

## Schema

| Column | Type | Description |
|--------|------|-------------|
| `id` | string | Unique test case identifier |
| `prompt` | string | Test prompt/question |
| `expected_tool` | string | Expected tool to be called |
| `expected_answer` | string | Expected answer (if applicable) |
| `difficulty` | string | easy, medium, or hard |
| `category` | string | Test category |
| `agent_type` | string | Suitable agent type |

## Usage

```bash
# Quick evaluation (default dataset)
smoltrace-eval --model gpt-4

# Explicit dataset specification
smoltrace-eval --model claude-3-sonnet --dataset-name {username}/smoltrace-tasks
```

```python
from datasets import load_dataset

# Load tasks
ds = load_dataset("{username}/smoltrace-tasks")

print(f"Total tasks: {{len(ds['train'])}}")
```

## Related Datasets

- **[smoltrace-benchmark-v1]({SMOLTRACE_REPO_URL})**: Full benchmark (132 cases)
- **smoltrace-leaderboard**: Aggregated results
{_get_footer()}"""
