app/routers/training.py

"""
Training API Router - Endpoints for model training
Enhanced with dataset column mapping and prompt structure controls
"""

from fastapi import APIRouter, HTTPException, BackgroundTasks, Depends
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field, validator
from typing import Optional, Dict, List, Any, Union
from datetime import datetime
import uuid
import logging
import re

from app.services.queue_service import TrainingQueue, JobPriority
from app.database import get_db, TrainingJob, JobStatus, TaskType
from app.config import TRAINING_TEMPLATES, settings
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select

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


# ============================================
# SIMPLIFIED REQUEST MODELS (matching dashboard form)
# ============================================

class DatasetConfigSimple(BaseModel):
    """Simplified dataset config matching the dashboard form."""
    name: str = Field(..., description="HuggingFace dataset name")
    train_split: str = Field(default="train")
    validation_split: Optional[str] = Field(default="validation")
    column_mapping: Dict[str, str] = Field(default_factory=dict, description="Maps roles to column names: {text: 'col1', input: 'col2'}")
    max_length: int = Field(default=512)


class TrainingArgsSimple(BaseModel):
    """Simplified training args matching the dashboard form."""
    epochs: int = Field(default=3)
    batch_size: int = Field(default=1)
    learning_rate: float = Field(default=5e-5)
    warmup_steps: int = Field(default=100)


class PEFTConfigSimple(BaseModel):
    """Simplified PEFT config matching the dashboard form."""
    enabled: bool = Field(default=True)
    method: str = Field(default="lora")
    r: int = Field(default=16)
    alpha: int = Field(default=32)
    dropout: float = Field(default=0.05)


class PromptTemplateSimple(BaseModel):
    """Simplified prompt template matching the dashboard form."""
    preset: str = Field(default="none", description="Template preset: none, alpaca, chatml, llama3, mistral, vicuna, phi3, reasoning")
    custom: Optional[Dict[str, Any]] = Field(default=None, description="Custom template sections")


class TrainingRequestSimple(BaseModel):
    """Simplified training request matching the dashboard form."""
    name: str = Field(default="training-job")
    task_type: str = Field(default="causal-lm")
    base_model: str = Field(..., description="HuggingFace model ID")
    dataset: DatasetConfigSimple
    training_args: TrainingArgsSimple = Field(default_factory=TrainingArgsSimple)
    peft_config: Optional[PEFTConfigSimple] = Field(None)
    prompt_template: Optional[PromptTemplateSimple] = Field(None, description="Prompt template configuration")


class TrainingJobResponse(BaseModel):
    """Training job response."""
    job_id: str
    status: str
    message: str
    created_at: datetime


class JobStatusResponse(BaseModel):
    """Job status response."""
    job_id: str
    name: str
    status: str
    progress: float
    current_epoch: int
    total_epochs: int
    current_step: int
    total_steps: int
    train_loss: Optional[float]
    eval_loss: Optional[float]
    metrics: Dict[str, Any]
    error_message: Optional[str]
    created_at: datetime
    started_at: Optional[datetime]
    completed_at: Optional[datetime]
    output_path: Optional[str]
    hub_model_id: Optional[str]


class DatasetPreviewResponse(BaseModel):
    """Dataset preview with detected columns and sample data."""
    dataset_name: str
    config: Optional[str]
    splits: List[str]
    columns: List[str]
    sample_data: List[Dict[str, Any]]
    total_rows: Optional[int]
    detected_task_types: List[str]
    suggested_column_mapping: Dict[str, str]


# Global queue instance
training_queue = None


def get_queue():
    """Get or create training queue."""
    global training_queue
    if training_queue is None:
        training_queue = TrainingQueue(max_concurrent=settings.MAX_CONCURRENT_JOBS)
    return training_queue


# ============================================
# DATASET PREVIEW AND COLUMN DETECTION
# ============================================

@router.get("/dataset/preview/{dataset_name:path}", response_model=DatasetPreviewResponse)
async def preview_dataset(
    dataset_name: str,
    config: Optional[str] = None,
    split: str = "train",
    rows: int = 10
):
    """Preview a dataset with detected columns and suggested mappings."""
    try:
        from datasets import load_dataset
        
        # Load dataset
        if config:
            ds = load_dataset(dataset_name, config, split=split, trust_remote_code=True, streaming=False)
        else:
            ds = load_dataset(dataset_name, split=split, trust_remote_code=True, streaming=False)
        
        # Get column info
        columns = []
        for col_name, col_type in ds.features.items():
            columns.append(col_name)
        
        # Get sample data
        sample_data = []
        for i in range(min(rows, len(ds))):
            sample_data.append({k: str(v)[:100] if v else None for k, v in ds[i].items()})
        
        # Detect task type and suggest column mapping
        detected_tasks, suggested_mapping = detect_task_and_mapping(ds)
        
        # Get all splits
        try:
            from datasets import load_dataset_builder
            builder = load_dataset_builder(dataset_name, trust_remote_code=True)
            splits = list(builder.info.splits.keys())
        except:
            splits = [split]
        
        return DatasetPreviewResponse(
            dataset_name=dataset_name,
            config=config,
            splits=splits,
            columns=columns,
            sample_data=sample_data,
            total_rows=len(ds),
            detected_task_types=detected_tasks,
            suggested_column_mapping=suggested_mapping
        )
        
    except Exception as e:
        logger.error(f"Error loading dataset: {e}")
        raise HTTPException(status_code=400, detail=f"Error loading dataset: {str(e)}")


def detect_task_and_mapping(dataset) -> tuple:
    """Detect suitable task types and suggest column mappings."""
    col_names_lower = [c.lower() for c in dataset.column_names]
    col_names_original = list(dataset.column_names)
    detected_tasks = []
    mapping = {}
    
    # Build a mapping from lowercase to original
    col_map = {c.lower(): c for c in col_names_original}
    
    # Check for common patterns
    
    # Text classification
    if 'label' in col_names_lower and 'text' in col_names_lower:
        detected_tasks.append("text-classification")
        mapping['label'] = col_map['label']
        mapping['text'] = col_map['text']
    
    # QA
    if 'question' in col_names_lower and 'answer' in col_names_lower:
        detected_tasks.append("question-answering")
        mapping['question'] = col_map['question']
        mapping['answer'] = col_map.get('answer', col_map.get('answers', ''))
        if 'context' in col_names_lower:
            mapping['context'] = col_map['context']
    
    # Instruction-output
    if 'instruction' in col_names_lower:
        detected_tasks.append("causal-lm")
        mapping['instruction'] = col_map['instruction']
        if 'input' in col_names_lower:
            mapping['input'] = col_map['input']
        if 'output' in col_names_lower:
            mapping['output'] = col_map['output']
    
    # Input-output
    if 'input' in col_names_lower and 'output' in col_names_lower:
        if 'causal-lm' not in detected_tasks:
            detected_tasks.append("causal-lm")
        mapping['input'] = col_map['input']
        mapping['output'] = col_map['output']
    
    # Reasoning
    if 'reasoning' in col_names_lower or 'thinking' in col_names_lower:
        detected_tasks.append("reasoning")
        if 'reasoning' in col_names_lower:
            mapping['reasoning'] = col_map['reasoning']
    
    # Default
    if not detected_tasks:
        detected_tasks.append("causal-lm")
        # Use first text-like column
        for col in col_names_original:
            if len(dataset) > 0 and isinstance(dataset[0].get(col), str):
                mapping['text'] = col
                break
    
    return detected_tasks, mapping


# ============================================
# TRAINING JOB ENDPOINTS
# ============================================

@router.post("/start", response_model=TrainingJobResponse)
async def start_training(
    request: TrainingRequestSimple,
    db: AsyncSession = Depends(get_db)
):
    """Start a new training job."""
    queue = get_queue()
    job_id = str(uuid.uuid4())
    
    # Create database record
    training_job = TrainingJob(
        job_id=job_id,
        name=request.name,
        task_type=request.task_type,
        base_model=request.base_model,
        dataset_name=request.dataset.name,
        dataset_split=request.dataset.train_split,
        training_args={
            "epochs": request.training_args.epochs,
            "batch_size": request.training_args.batch_size,
            "learning_rate": request.training_args.learning_rate,
            "warmup_steps": request.training_args.warmup_steps,
        },
        peft_config=request.peft_config.dict() if request.peft_config else None,
        status=JobStatus.PENDING.value,
        total_epochs=request.training_args.epochs,
    )
    
    db.add(training_job)
    await db.commit()
    
    # Build full config for training service
    config = {
        "job_id": job_id,
        "task_type": request.task_type,
        "base_model": request.base_model,
        "model_name": request.base_model,
        "dataset_name": request.dataset.name,
        "dataset": {
            "name": request.dataset.name,
            "train_split": request.dataset.train_split,
            "validation_split": request.dataset.validation_split,
            "column_mapping": request.dataset.column_mapping,
            "max_length": request.dataset.max_length,
        },
        "training_args": {
            "epochs": request.training_args.epochs,
            "batch_size": request.training_args.batch_size,
            "learning_rate": request.training_args.learning_rate,
            "warmup_steps": request.training_args.warmup_steps,
        },
        "peft_config": request.peft_config.dict() if request.peft_config else None,
        "prompt_template": request.prompt_template.dict() if request.prompt_template else {"preset": "none"},
    }
    
    # Submit to queue
    priority = JobPriority.NORMAL
    await queue.submit(config, priority=priority)
    
    # Update status
    training_job.status = JobStatus.QUEUED.value
    await db.commit()
    
    logger.info(f"Training job {job_id} submitted: {request.name}")
    
    return TrainingJobResponse(
        job_id=job_id,
        status="queued",
        message="Training job submitted successfully",
        created_at=datetime.utcnow()
    )


@router.get("/status/{job_id}", response_model=JobStatusResponse)
async def get_job_status(
    job_id: str,
    db: AsyncSession = Depends(get_db)
):
    """Get status of a training job."""
    result = await db.execute(
        select(TrainingJob).where(TrainingJob.job_id == job_id)
    )
    job = result.scalar_one_or_none()
    
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    
    return JobStatusResponse(
        job_id=job.job_id,
        name=job.name,
        status=job.status,
        progress=job.progress or 0.0,
        current_epoch=job.current_epoch or 0,
        total_epochs=job.total_epochs or 0,
        current_step=job.current_step or 0,
        total_steps=job.total_steps or 0,
        train_loss=job.train_loss,
        eval_loss=job.eval_loss,
        metrics=job.metrics or {},
        error_message=job.error_message,
        created_at=job.created_at,
        started_at=job.started_at,
        completed_at=job.completed_at,
        output_path=job.output_path,
        hub_model_id=job.hub_model_id
    )


@router.get("/jobs")
async def list_jobs(
    status: Optional[str] = None,
    limit: int = 50,
    offset: int = 0,
    db: AsyncSession = Depends(get_db)
):
    """List all training jobs."""
    query = select(TrainingJob).order_by(TrainingJob.created_at.desc())
    
    if status:
        query = query.where(TrainingJob.status == status)
    
    query = query.offset(offset).limit(limit)
    
    result = await db.execute(query)
    jobs = result.scalars().all()
    
    return {
        "jobs": [
            {
                "job_id": job.job_id,
                "name": job.name,
                "status": job.status,
                "progress": job.progress or 0.0,
                "current_epoch": job.current_epoch or 0,
                "total_epochs": job.total_epochs or 0,
                "current_step": job.current_step or 0,
                "total_steps": job.total_steps or 0,
                "train_loss": job.train_loss,
                "eval_loss": job.eval_loss,
                "metrics": job.metrics or {},
                "error_message": job.error_message,
                "created_at": job.created_at.isoformat() if job.created_at else None,
                "started_at": job.started_at.isoformat() if job.started_at else None,
                "completed_at": job.completed_at.isoformat() if job.completed_at else None,
                "model_name": job.base_model,
                "dataset_name": job.dataset_name,
            }
            for job in jobs
        ]
    }


@router.post("/cancel/{job_id}")
async def cancel_job(
    job_id: str,
    db: AsyncSession = Depends(get_db)
):
    """Cancel a training job."""
    queue = get_queue()
    cancelled = await queue.cancel_job(job_id)
    
    result = await db.execute(
        select(TrainingJob).where(TrainingJob.job_id == job_id)
    )
    job = result.scalar_one_or_none()
    
    if job:
        job.status = JobStatus.CANCELLED.value
        await db.commit()
    
    return {"message": f"Job {job_id} cancelled", "success": True}


@router.get("/queue/status")
async def get_queue_status():
    """Get current queue status."""
    queue = get_queue()
    return {
        "queue_size": await queue.get_queue_size(),
        "active_jobs": await queue.get_active_jobs(),
        "max_concurrent": settings.MAX_CONCURRENT_JOBS
    }


@router.get("/templates")
async def get_training_templates():
    """Get available training configuration templates."""
    return TRAINING_TEMPLATES


@router.get("/metrics/{job_id}")
async def get_job_metrics(
    job_id: str,
    db: AsyncSession = Depends(get_db)
):
    """Get detailed metrics for a training job."""
    result = await db.execute(
        select(TrainingJob).where(TrainingJob.job_id == job_id)
    )
    job = result.scalar_one_or_none()
    
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    
    return {
        "job_id": job_id,
        "train_loss": job.train_loss,
        "eval_loss": job.eval_loss,
        "metrics": job.metrics
    }


@router.delete("/job/{job_id}")
async def delete_job(
    job_id: str,
    db: AsyncSession = Depends(get_db)
):
    """Delete a completed training job record."""
    result = await db.execute(
        select(TrainingJob).where(TrainingJob.job_id == job_id)
    )
    job = result.scalar_one_or_none()
    
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    
    if job.status not in [JobStatus.COMPLETED.value, JobStatus.FAILED.value, JobStatus.CANCELLED.value]:
        raise HTTPException(status_code=400, detail="Cannot delete active job")
    
    await db.delete(job)
    await db.commit()
    
    return {"message": f"Job {job_id} deleted", "success": True}


# ============================================
# PROMPT TEMPLATE ENDPOINTS
# ============================================

@router.get("/prompt-templates")
async def get_prompt_templates():
    """Get available prompt template presets."""
    return {
        "presets": [
            {"id": "none", "name": "None (Raw Text)", "description": "Use dataset text directly"},
            {"id": "alpaca", "name": "Alpaca Format", "description": "Instruction-Input-Output"},
            {"id": "chatml", "name": "ChatML", "description": "ChatML format"},
            {"id": "llama3", "name": "Llama 3", "description": "Llama 3 instruction format"},
            {"id": "mistral", "name": "Mistral", "description": "Mistral instruction format"},
            {"id": "vicuna", "name": "Vicuna", "description": "Vicuna chat format"},
            {"id": "phi3", "name": "Phi-3", "description": "Microsoft Phi-3 format"},
            {"id": "reasoning", "name": "Reasoning/CoT", "description": "Chain-of-thought"}
        ]
    }