app/routers/models.py

"""
Models API Router - Explore and select models 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/models", tags=["Models"])


class ModelInfo(BaseModel):
    """Model information."""
    id: str
    author: str
    model_name: str
    pipeline_tag: Optional[str]
    library_name: Optional[str]
    downloads: int = 0
    likes: int = 0
    private: bool = False
    tags: List[str] = []
    created_at: Optional[datetime]
    last_modified: Optional[datetime]
    card_data: Optional[Dict[str, Any]]
    

class ModelSearchResult(BaseModel):
    """Search result model."""
    id: str
    author: str
    model_name: str
    pipeline_tag: Optional[str]
    downloads: int
    likes: int
    

class ModelRecommendation(BaseModel):
    """Model recommendation for a task."""
    model_id: str
    reason: str
    estimated_vram: str
    suggested_batch_size: int
    supports_peft: bool
    

class ModelCompatibility(BaseModel):
    """Model compatibility check result."""
    model_id: str
    compatible: bool
    issues: List[str]
    recommendations: List[str]
    max_sequence_length: int
    estimated_training_time: str
    requires_gpu: bool


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


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


@router.get("/info/{model_id:path}", response_model=ModelInfo)
async def get_model_info(model_id: str):
    """Get detailed information about a model."""
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(
                f"{HF_API_URL}/models/{model_id}",
                timeout=30.0
            )
            response.raise_for_status()
            data = response.json()
            
            return ModelInfo(
                id=data.get("id", ""),
                author=data.get("author", ""),
                model_name=data.get("id", "").split("/")[-1] if "/" in data.get("id", "") else data.get("id", ""),
                pipeline_tag=data.get("pipeline_tag"),
                library_name=data.get("library_name"),
                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,
                card_data=data.get("cardData")
            )
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 404:
                raise HTTPException(status_code=404, detail="Model 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 model info: {e}")
            raise HTTPException(status_code=500, detail=str(e))


@router.get("/recommend/{task_type}", response_model=List[ModelRecommendation])
async def get_model_recommendations(task_type: str):
    """Get model recommendations for a specific task."""
    recommendations = {
        "causal-lm": [
            ModelRecommendation(model_id="microsoft/phi-2", reason="Compact but powerful, excellent for fine-tuning", estimated_vram="8GB", suggested_batch_size=4, supports_peft=True),
            ModelRecommendation(model_id="google/gemma-2b", reason="Efficient 2B model, great for instruction tuning", estimated_vram="6GB", suggested_batch_size=8, supports_peft=True),
            ModelRecommendation(model_id="mistralai/Mistral-7B-v0.1", reason="Best-in-class 7B model", estimated_vram="16GB", suggested_batch_size=2, supports_peft=True),
            ModelRecommendation(model_id="meta-llama/Llama-3.2-3B", reason="Latest Llama architecture", estimated_vram="10GB", suggested_batch_size=4, supports_peft=True),
        ],
        "seq2seq": [
            ModelRecommendation(model_id="google/flan-t5-base", reason="Versatile instruction-following model", estimated_vram="4GB", suggested_batch_size=16, supports_peft=True),
            ModelRecommendation(model_id="google/flan-t5-large", reason="Larger variant for better quality", estimated_vram="8GB", suggested_batch_size=8, supports_peft=True),
            ModelRecommendation(model_id="facebook/bart-large", reason="Strong summarization model", estimated_vram="6GB", suggested_batch_size=8, supports_peft=True),
        ],
        "summarization": [
            ModelRecommendation(model_id="facebook/bart-large-cnn", reason="Pre-trained for news summarization", estimated_vram="6GB", suggested_batch_size=8, supports_peft=True),
            ModelRecommendation(model_id="google/flan-t5-base", reason="Multi-purpose, handles various summarization tasks", estimated_vram="4GB", suggested_batch_size=16, supports_peft=True),
            ModelRecommendation(model_id="mistralai/Mistral-7B-v0.1", reason="Large language model for advanced summarization", estimated_vram="16GB", suggested_batch_size=2, supports_peft=True),
        ],
        "token-classification": [
            ModelRecommendation(model_id="dslim/bert-base-NER", reason="Strong NER model, 18 entity types", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True),
            ModelRecommendation(model_id="dbmdz/bert-large-cased-finetuned-conll03-english", reason="Large NER model for high accuracy", estimated_vram="4GB", suggested_batch_size=16, supports_peft=True),
            ModelRecommendation(model_id="distilbert-base-cased", reason="Fast and efficient for fine-tuning", estimated_vram="1GB", suggested_batch_size=64, supports_peft=True),
        ],
        "sequence-classification": [
            ModelRecommendation(model_id="distilbert-base-uncased-finetuned-sst-2-english", reason="Sentiment analysis baseline", estimated_vram="1GB", suggested_batch_size=64, supports_peft=True),
            ModelRecommendation(model_id="roberta-base", reason="Strong general-purpose classification model", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True),
            ModelRecommendation(model_id="microsoft/deberta-v3-base", reason="State-of-the-art classification accuracy", estimated_vram="4GB", suggested_batch_size=16, supports_peft=True),
        ],
        "question-answering": [
            ModelRecommendation(model_id="deepset/roberta-base-squad2", reason="SQuAD2.0 trained, handles unanswerable", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True),
            ModelRecommendation(model_id="distilbert-base-cased-distilled-squad", reason="Fast QA model", estimated_vram="1GB", suggested_batch_size=64, supports_peft=True),
            ModelRecommendation(model_id="bert-large-uncased-whole-word-masking-finetuned-squad", reason="High accuracy QA", estimated_vram="4GB", suggested_batch_size=16, supports_peft=True),
        ],
        "translation": [
            ModelRecommendation(model_id="Helsinki-NLP/opus-mt-en-de", reason="English to German", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True),
            ModelRecommendation(model_id="Helsinki-NLP/opus-mt-en-fr", reason="English to French", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True),
            ModelRecommendation(model_id="facebook/mbart-large-50-many-to-many-mmt", reason="Multilingual translation", estimated_vram="8GB", suggested_batch_size=8, supports_peft=True),
        ],
        "masked-lm": [
            ModelRecommendation(model_id="bert-base-uncased", reason="Classic MLM model", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True),
            ModelRecommendation(model_id="roberta-base", reason="Modern MLM architecture", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True),
            ModelRecommendation(model_id="distilbert-base-uncased", reason="Fast MLM training", estimated_vram="1GB", suggested_batch_size=64, supports_peft=True),
        ],
    }
    
    result = recommendations.get(task_type, [])
    
    if not result:
        result = recommendations.get("causal-lm", [])
    
    return result[:5]


@router.get("/check/{model_id:path}", response_model=ModelCompatibility)
async def check_model_compatibility(
    model_id: str,
    task: str = Query("causal-lm", description="Target task")
):
    """Check if a model is compatible for training with current setup."""
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(
                f"{HF_API_URL}/models/{model_id}",
                timeout=30.0
            )
            response.raise_for_status()
            data = response.json()
            
            issues = []
            recommendations = []
            compatible = True
            
            library = data.get("library_name", "transformers")
            if library not in ["transformers", "peft", "adapter-transformers"]:
                issues.append(f"Unsupported library: {library}")
                compatible = False
            
            tags = data.get("tags", [])
            
            if "gated" in tags:
                recommendations.append("This is a gated model. Ensure you have access and set HF_TOKEN.")
            
            if data.get("private", False):
                recommendations.append("Private model requires authentication.")
            
            pipeline_tag = data.get("pipeline_tag", "")
            
            if "bert" in model_id.lower():
                max_seq_len = 512
            elif "llama" in model_id.lower():
                max_seq_len = 4096
            elif "mistral" in model_id.lower():
                max_seq_len = 32768
            elif "t5" in model_id.lower():
                max_seq_len = 512
            else:
                max_seq_len = 2048
            
            task_pipeline_map = {
                "causal-lm": "text-generation",
                "seq2seq": "text2text-generation",
                "summarization": "summarization",
                "translation": "translation",
                "token-classification": "token-classification",
                "sequence-classification": "text-classification",
                "question-answering": "question-answering",
                "masked-lm": "fill-mask"
            }
            
            expected_pipeline = task_pipeline_map.get(task)
            if expected_pipeline and pipeline_tag and pipeline_tag != expected_pipeline:
                recommendations.append(f"Model pipeline tag ({pipeline_tag}) may not match task ({task}).")
            
            requires_gpu = False
            model_name_lower = model_id.lower()
            if any(x in model_name_lower for x in ["7b", "13b", "30b", "70b"]):
                issues.append("Large models (7B+) require GPU for practical training")
                requires_gpu = True
                compatible = False
            
            if "7b" in model_id.lower():
                training_time = "Hours per epoch on GPU"
            elif "2b" in model_id.lower() or "3b" in model_id.lower():
                training_time = "30-60 min per epoch on GPU"
            else:
                training_time = "10-30 min per epoch on GPU"
            
            return ModelCompatibility(
                model_id=model_id,
                compatible=compatible,
                issues=issues,
                recommendations=recommendations,
                max_sequence_length=max_seq_len,
                estimated_training_time=training_time,
                requires_gpu=requires_gpu
            )
            
        except httpx.HTTPStatusError as e:
            if e.response.status_code == 404:
                raise HTTPException(status_code=404, detail="Model not found")
            raise HTTPException(status_code=e.response.status_code, detail=f"API error")
        except Exception as e:
            logger.error(f"Compatibility check failed: {e}")
            raise HTTPException(status_code=500, detail=str(e))


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


@router.get("/tasks")
async def get_supported_tasks():
    """Get list of supported training tasks."""
    return {
        "tasks": [
            {"id": "causal-lm", "name": "Causal Language Modeling", "description": "Text generation, instruction tuning, chat models"},
            {"id": "seq2seq", "name": "Sequence-to-Sequence", "description": "Input-output transformations"},
            {"id": "summarization", "name": "Summarization", "description": "Text summarization models"},
            {"id": "translation", "name": "Translation", "description": "Language translation models"},
            {"id": "token-classification", "name": "Token Classification", "description": "NER, POS tagging, entity extraction"},
            {"id": "sequence-classification", "name": "Sequence Classification", "description": "Sentiment, topic, intent classification"},
            {"id": "question-answering", "name": "Question Answering", "description": "Extractive QA models"},
            {"id": "masked-lm", "name": "Masked Language Modeling", "description": "BERT-style pretraining"},
        ]
    }