app/routers/datasets.py

"""
Datasets API Router - Explore and select datasets from HuggingFace Hub
"""

from fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel, Field
from typing import Optional, Dict, List, Any
from datetime import datetime
import logging
import httpx

from app.config import settings

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/datasets", tags=["Datasets"])


class DatasetInfo(BaseModel):
    """Dataset information."""
    id: str
    author: str
    dataset_name: str
    description: Optional[str]
    downloads: int = 0
    likes: int = 0
    private: bool = False
    tags: List[str] = []
    created_at: Optional[datetime]
    last_modified: Optional[datetime]
    

class DatasetSearchResult(BaseModel):
    """Search result dataset."""
    id: str
    author: str
    dataset_name: str
    downloads: int
    likes: int
    

class DatasetRecommendation(BaseModel):
    """Dataset recommendation for a task."""
    dataset_id: str
    reason: str
    size: str
    format: str


class DatasetSchema(BaseModel):
    """Dataset schema information."""
    dataset_id: str
    splits: List[str]
    columns: List[str]
    num_rows: Dict[str, int]
    features: Dict[str, str]


HF_API_URL = "https://huggingface.co/api"


@router.get("/search", response_model=List[DatasetSearchResult])
async def search_datasets(
    query: str = Query(..., min_length=1, description="Search query"),
    task: Optional[str] = Query(None, description="Filter by task"),
    language: Optional[str] = Query(None, description="Filter by language"),
    author: Optional[str] = Query(None, description="Filter by author"),
    limit: int = Query(20, ge=1, le=100)
):
    """Search HuggingFace Hub for datasets."""
    async with httpx.AsyncClient() as client:
        params = {
            "search": query,
            "limit": limit,
            "full": "false"
        }
        
        if task:
            params["task_categories"] = task
        if language:
            params["language"] = language
        if author:
            params["author"] = author
        
        try:
            response = await client.get(
                f"{HF_API_URL}/datasets",
                params=params,
                timeout=30.0
            )
            response.raise_for_status()
            data = response.json()
            
            return [
                DatasetSearchResult(
                    id=d.get("id", ""),
                    author=d.get("author", ""),
                    dataset_name=d.get("id", "").split("/")[-1] if "/" in d.get("id", "") else d.get("id", ""),
                    downloads=d.get("downloads", 0),
                    likes=d.get("likes", 0)
                )
                for d in data
            ]
        except Exception as e:
            logger.error(f"Dataset search failed: {e}")
            raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}")


@router.get("/info/{dataset_id:path}", response_model=DatasetInfo)
async def get_dataset_info(dataset_id: str):
    """Get detailed information about a dataset."""
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(
                f"{HF_API_URL}/datasets/{dataset_id}",
                timeout=30.0
            )
            response.raise_for_status()
            data = response.json()
            
            return DatasetInfo(
                id=data.get("id", ""),
                author=data.get("author", ""),
                dataset_name=data.get("id", "").split("/")[-1] if "/" in data.get("id", "") else data.get("id", ""),
                description=data.get("cardData", {}).get("description", ""),
                downloads=data.get("downloads", 0),
                likes=data.get("likes", 0),
                private=data.get("private", False),
                tags=data.get("tags", []),
                created_at=datetime.fromisoformat(data["createdAt"].replace("Z", "+00:00")) if data.get("createdAt") else None,
                last_modified=datetime.fromisoformat(data["lastModified"].replace("Z", "+00:00")) if data.get("lastModified") else None
            )
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 404:
                raise HTTPException(status_code=404, detail="Dataset not found")
            raise HTTPException(status_code=e.response.status_code, detail=f"API error: {e}")
        except Exception as e:
            logger.error(f"Failed to get dataset info: {e}")
            raise HTTPException(status_code=500, detail=str(e))


@router.get("/schema/{dataset_id:path}")
async def get_dataset_schema(dataset_id: str):
    """Get dataset schema including splits, columns, and row counts."""
    try:
        from datasets import get_dataset_infos
        
        infos = get_dataset_infos(dataset_id)
        if not infos:
            raise HTTPException(status_code=404, detail="Dataset not found or inaccessible")
        
        info = list(infos.values())[0]
        
        splits = list(info.splits.keys()) if info.splits else []
        columns = list(info.features.keys()) if info.features else []
        num_rows = {k: v.num_examples for k, v in info.splits.items()} if info.splits else {}
        features = {k: str(v) for k, v in info.features.items()} if info.features else {}
        
        return {
            "dataset_id": dataset_id,
            "splits": splits,
            "columns": columns,
            "num_rows": num_rows,
            "features": features
        }
    except Exception as e:
        logger.error(f"Failed to get dataset schema: {e}")
        return {
            "dataset_id": dataset_id,
            "splits": [],
            "columns": [],
            "num_rows": {},
            "features": {},
            "error": str(e)
        }


@router.get("/recommend/{task_type}", response_model=List[DatasetRecommendation])
async def get_dataset_recommendations(task_type: str):
    """Get dataset recommendations for a specific task."""
    recommendations = {
        "causal-lm": [
            DatasetRecommendation(dataset_id="tiiuae/falcon-refinedweb", reason="Large-scale web data for language modeling", size="350GB+", format="Raw text"),
            DatasetRecommendation(dataset_id="OpenAssistant/oasst1", reason="Human conversations for instruction tuning", size="~10M tokens", format="Conversations"),
            DatasetRecommendation(dataset_id="HuggingFaceH4/ultrachat_200k", reason="High-quality dialogues", size="200k conversations", format="Conversations"),
            DatasetRecommendation(dataset_id="teknium/openhermes", reason="Diverse instruction data", size="1M+ samples", format="Instructions"),
        ],
        "seq2seq": [
            DatasetRecommendation(dataset_id="cfilt/iitb-english-hindi", reason="Translation dataset", size="~1.5M pairs", format="Parallel text"),
            DatasetRecommendation(dataset_id="wmt/wmt14-en-de", reason="WMT translation benchmark", size="~4.5M pairs", format="Parallel text"),
        ],
        "summarization": [
            DatasetRecommendation(dataset_id="cnn_dailymail", reason="News summarization standard", size="~300k articles", format="Article-Summary"),
            DatasetRecommendation(dataset_id="xsum", reason="BBC article summaries", size="~22k articles", format="Article-Summary"),
            DatasetRecommendation(dataset_id="samsum", reason="Dialogue summarization", size="~16k dialogues", format="Dialogue-Summary"),
        ],
        "token-classification": [
            DatasetRecommendation(dataset_id="conll2003", reason="Classic NER benchmark", size="~22k sentences", format="IOB tags"),
            DatasetRecommendation(dataset_id="ontonotes", reason="Multi-domain NER", size="~80k sentences", format="IOB tags"),
            DatasetRecommendation(dataset_id="wikiann", reason="Cross-lingual NER", size="Varies by language", format="IOB tags"),
        ],
        "sequence-classification": [
            DatasetRecommendation(dataset_id="imdb", reason="Movie reviews for sentiment", size="50k reviews", format="Text-Label"),
            DatasetRecommendation(dataset_id="yelp_polarity", reason="Yelp reviews", size="560k reviews", format="Text-Label"),
            DatasetRecommendation(dataset_id="ag_news", reason="News topic classification", size="~127k articles", format="Text-Label"),
        ],
        "question-answering": [
            DatasetRecommendation(dataset_id="squad", reason="Reading comprehension QA", size="~100k questions", format="Context-QA"),
            DatasetRecommendation(dataset_id="squad_v2", reason="SQuAD with unanswerable", size="~150k questions", format="Context-QA"),
            DatasetRecommendation(dataset_id="natural_questions", reason="Real user questions", size="~300k questions", format="Document-QA"),
        ],
        "translation": [
            DatasetRecommendation(dataset_id="wmt/wmt14-en-de", reason="English-German", size="~4.5M pairs", format="Parallel"),
            DatasetRecommendation(dataset_id="wmt/wmt19-ru-en", reason="Russian-English", size="~4M pairs", format="Parallel"),
        ],
        "masked-lm": [
            DatasetRecommendation(dataset_id="bookcorpus", reason="Books dataset for MLM", size="~1B words", format="Raw text"),
            DatasetRecommendation(dataset_id="wikitext", reason="Wikipedia text", size="Varies", format="Raw text"),
        ],
    }
    
    result = recommendations.get(task_type, [])
    
    if not result:
        result = recommendations.get("causal-lm", [])
    
    return result[:5]


@router.get("/popular")
async def get_popular_datasets(
    task: Optional[str] = None,
    limit: int = Query(10, ge=1, le=50)
):
    """Get most popular datasets, optionally filtered by task."""
    async with httpx.AsyncClient() as client:
        params = {
            "sort": "downloads",
            "direction": "-1",
            "limit": limit
        }
        
        if task:
            params["task_categories"] = task
        
        try:
            response = await client.get(
                f"{HF_API_URL}/datasets",
                params=params,
                timeout=30.0
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"Failed to get popular datasets: {e}")
            raise HTTPException(status_code=500, detail=str(e))


@router.get("/preview/{dataset_id:path}")
async def preview_dataset(
    dataset_id: str,
    split: str = Query("train", description="Dataset split to preview"),
    rows: int = Query(10, ge=1, le=100)
):
    """Preview first N rows of a dataset."""
    try:
        from datasets import load_dataset
        
        ds = load_dataset(dataset_id, split=split, streaming=True)
        data = []
        for i, row in enumerate(ds):
            if i >= rows:
                break
            data.append(row)
        
        return {
            "dataset_id": dataset_id,
            "split": split,
            "rows": data,
            "columns": list(data[0].keys()) if data else []
        }
    except Exception as e:
        logger.error(f"Failed to preview dataset: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/configs/{dataset_id:path}")
async def get_dataset_configs(dataset_id: str):
    """Get available configurations for a dataset."""
    try:
        from datasets import get_dataset_config_names
        
        configs = get_dataset_config_names(dataset_id)
        return {
            "dataset_id": dataset_id,
            "configs": configs
        }
    except Exception as e:
        logger.error(f"Failed to get configs: {e}")
        return {
            "dataset_id": dataset_id,
            "configs": [],
            "error": str(e)
        }


@router.get("/splits/{dataset_id:path}")
async def get_dataset_splits(
    dataset_id: str,
    config: Optional[str] = None
):
    """Get available splits for a dataset."""
    try:
        from datasets import get_dataset_split_names
        
        splits = get_dataset_split_names(dataset_id, config_name=config)
        return {
            "dataset_id": dataset_id,
            "config": config,
            "splits": splits
        }
    except Exception as e:
        logger.error(f"Failed to get splits: {e}")
        return {
            "dataset_id": dataset_id,
            "config": config,
            "splits": [],
            "error": str(e)
        }