#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "transformers>=4.40.0",
#     "torch>=2.0.0",
#     "rich>=13.0.0",
# ]
# ///
"""
Example uv script with inline dependencies.

This script demonstrates the PEP 723 inline metadata format for uv scripts.
Dependencies are declared in the script header and installed automatically.

Usage:
    # Run locally
    uv run example_uv_script.py

    # Run on remote HF Space via SSH
    scp example_uv_script.py user@SPACE.hf.space:/app/
    ssh user@SPACE.hf.space "cd /app && uv run example_uv_script.py"

    # Make executable and run directly
    chmod +x example_uv_script.py
    ./example_uv_script.py
"""

from rich.console import Console
from rich.table import Table
import torch


def check_cuda() -> dict:
    """Check CUDA availability and device info."""
    info = {
        "cuda_available": torch.cuda.is_available(),
        "device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0,
        "device_name": None,
        "cuda_version": None,
    }

    if info["cuda_available"]:
        info["device_name"] = torch.cuda.get_device_name(0)
        info["cuda_version"] = torch.version.cuda

    return info


def run_inference():
    """Run a simple inference example."""
    from transformers import pipeline

    # Use CPU if CUDA not available
    device = 0 if torch.cuda.is_available() else -1

    console = Console()
    console.print("\n[bold blue]Running sentiment analysis...[/bold blue]\n")

    classifier = pipeline("sentiment-analysis", device=device)
    texts = [
        "I love using Hugging Face!",
        "This is a challenging problem.",
        "The weather is nice today.",
    ]

    results = classifier(texts)

    # Display results in a table
    table = Table(title="Sentiment Analysis Results")
    table.add_column("Text", style="cyan")
    table.add_column("Label", style="green")
    table.add_column("Score", style="magenta")

    for text, result in zip(texts, results):
        table.add_row(text, result["label"], f"{result['score']:.4f}")

    console.print(table)
    return results


def main():
    console = Console()

    # Header
    console.print("\n[bold green]HF Space Sandbox - Example Script[/bold green]")
    console.print("=" * 50)

    # Check CUDA
    cuda_info = check_cuda()

    info_table = Table(title="System Information")
    info_table.add_column("Property", style="cyan")
    info_table.add_column("Value", style="yellow")

    info_table.add_row("CUDA Available", str(cuda_info["cuda_available"]))
    info_table.add_row("Device Count", str(cuda_info["device_count"]))
    if cuda_info["device_name"]:
        info_table.add_row("GPU", cuda_info["device_name"])
        info_table.add_row("CUDA Version", cuda_info["cuda_version"])

    console.print(info_table)

    # Run inference
    run_inference()

    console.print("\n[bold green]Done![/bold green]\n")


if __name__ == "__main__":
    main()