"""orbit badge — generate quality badges for HuggingFace dataset cards."""

from __future__ import annotations

from orbit.utils.display import console


GRADE_COLORS = {
    "A": "brightgreen",
    "B": "green",
    "C": "yellow",
    "D": "orange",
    "F": "red",
}


def generate_badge_markdown(grade: str, score: int | None = None) -> str:
    """Generate shields.io badge markdown for an ORBIT grade."""
    color = GRADE_COLORS.get(grade, "lightgrey")
    label = f"Grade%20{grade}"
    if score is not None:
        label += f"%20({score}/100)"
    return f"![ORBIT Grade: {grade}](https://img.shields.io/badge/ORBIT-{label}-{color})"


def run_badge(dataset_repo: str, push: bool = False) -> None:
    """Generate and optionally push an ORBIT badge to a dataset card."""
    import sys

    from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data
    from orbit.analyzer.quality import compute_quality
    from orbit.analyzer.success_predictor import DatasetReadinessGrader

    console.print(f"\n[bold]Generating ORBIT badge for {dataset_repo}...[/bold]")

    # Load and analyze
    try:
        info = load_dataset_info(dataset_repo)
        episode_data = load_episode_data(dataset_repo, max_episodes=50)
    except Exception as e:
        console.print(f"[bold red]Error:[/bold red] Could not load dataset: {e}")
        sys.exit(1)

    # Extract action episodes
    action_episodes = []
    for ep_idx in sorted(episode_data["episode_index"].unique()):
        ep = episode_data[episode_data["episode_index"] == ep_idx]
        actions = ep["action"].tolist()
        if actions:
            import numpy as np
            action_episodes.append(np.array([a for a in actions if a is not None]))

    if not action_episodes:
        console.print("[bold red]Error:[/bold red] No action data found.")
        sys.exit(1)

    # Grade
    grader = DatasetReadinessGrader()
    result = grader.grade(info, action_episodes)
    badge_md = generate_badge_markdown(result.grade, result.score)

    console.print(f"\n[bold green]Badge generated:[/bold green]")
    console.print(f"  Grade: [bold]{result.grade}[/bold] (score: {result.score}/100)")
    console.print(f"\n  [dim]Markdown:[/dim]")
    console.print(f"  {badge_md}")

    if push:
        _push_badge(dataset_repo, badge_md)
    else:
        console.print(f"\n[dim]Add --push to update the dataset card on HuggingFace.[/dim]")


def _push_badge(dataset_repo: str, badge_md: str) -> None:
    """Push badge to HuggingFace dataset card."""
    try:
        from huggingface_hub import HfApi

        api = HfApi()

        # Read existing README
        try:
            readme_content = api.hf_hub_download(
                repo_id=dataset_repo, filename="README.md", repo_type="dataset"
            )
            with open(readme_content) as f:
                content = f.read()
        except Exception:
            content = f"---\n---\n# {dataset_repo}\n"

        # Remove existing ORBIT badge if present
        lines = content.split("\n")
        lines = [l for l in lines if "img.shields.io/badge/ORBIT" not in l]

        # Insert badge after YAML frontmatter
        insert_idx = 0
        in_frontmatter = False
        for i, line in enumerate(lines):
            if line.strip() == "---":
                if not in_frontmatter:
                    in_frontmatter = True
                else:
                    insert_idx = i + 1
                    break

        lines.insert(insert_idx, "")
        lines.insert(insert_idx + 1, badge_md)
        lines.insert(insert_idx + 2, "")

        new_content = "\n".join(lines)

        api.upload_file(
            path_or_fileobj=new_content.encode(),
            path_in_repo="README.md",
            repo_id=dataset_repo,
            repo_type="dataset",
            commit_message="Add ORBIT quality badge",
        )
        console.print("[bold green]\u2713[/bold green] Badge pushed to dataset card!")

    except Exception as e:
        console.print(f"[bold red]Error:[/bold red] Could not push badge: {e}")
        console.print("[dim]Make sure you have write access to the dataset.[/dim]")
