# ta_ml/api/train.py
"""
API endpoints for model training.
"""
import logging
from typing import List, Tuple

import pandas as pd
from fastapi import (APIRouter, BackgroundTasks, Body, Depends, File,
                     HTTPException, UploadFile, status)

from api.dependencies import get_training_service
from data.file_handler import FileHandler, FileHandlerError
from models.api import (TrainDirectRequest, TrainingResponse)
from services.training_service import TrainingService

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


def _extract_data_from_df(df: pd.DataFrame) -> Tuple[List[str], List[str]]:
    """Helper to extract texts and labels from a DataFrame."""
    if "text" in df.columns and "label" in df.columns:
        return df["text"].dropna().astype(str).tolist(), df["label"].dropna().astype(str).tolist()

    # Fallback for wide format
    texts, labels = [], []
    label_columns = [col for col in df.columns if col.lower() not in [
        "nomor", "no"]]
    if not label_columns:
        raise HTTPException(
            status_code=400, detail="Could not find 'text'/'label' columns or any other valid label columns.")

    for _, row in df.iterrows():
        for col in label_columns:
            if pd.notna(row[col]):
                texts.append(str(row[col]))
                labels.append(col)

    if not texts:
        raise HTTPException(
            status_code=400, detail="No valid text-label pairs found in the file.")

    return texts, labels


def _run_training_task(
    training_service: TrainingService,
    texts: List[str],
    labels: List[str]
):
    """The actual training function to be run in the background."""
    try:
        logger.info(f"Background training started for {len(texts)} samples.")
        metrics = training_service.train_model(texts, labels)
        logger.info(
            f"Background training finished successfully. F1 Score: {metrics.f1_score:.4f}")
    except Exception as e:
        logger.error(f"Background training task failed: {e}", exc_info=True)


@router.post(
    "/file",
    response_model=TrainingResponse,
    summary="Train model from an uploaded file",
    status_code=status.HTTP_202_ACCEPTED,
    description=(
        "Initiates a model training process in the background using data from an uploaded file. "
        "The endpoint immediately returns a confirmation message while the training proceeds asynchronously. "
        "**Monitor the application logs to check for training completion and results.**\n\n"
        "The file can be in two formats:\n"
        "1.  **Long Format**: Must contain `text` and `label` columns.\n"
        "2.  **Wide Format**: Can have multiple columns where each column header is a label and the cells contain the corresponding texts."
    ),
    responses={
        202: {"description": "The training task has been accepted and started in the background."},
        400: {"description": "The uploaded file is invalid, has an unsupported format, or contains no valid data."},
        422: {"description": "Validation error, e.g., no file was uploaded."},
    },
)
async def train_from_file(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...,
                            description="A CSV or Excel file with training data."),
    training_service: TrainingService = Depends(get_training_service),
):
    try:
        df = FileHandler.read_to_dataframe(file)
        texts, labels = _extract_data_from_df(df)
    except FileHandlerError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except HTTPException as e:
        raise e

    background_tasks.add_task(
        _run_training_task, training_service, texts, labels)

    return TrainingResponse(message="Training has been accepted and started in the background. Check logs for completion status.")


@router.post(
    "/direct",
    response_model=TrainingResponse,
    summary="Train model from a direct JSON payload",
    status_code=status.HTTP_202_ACCEPTED,
    description=(
        "Initiates a model training process in the background using lists of texts and labels provided directly in the request body. "
        "The endpoint immediately returns a confirmation message. "
        "**Monitor the application logs to check for training completion and results.**"
    ),
    responses={
        202: {"description": "The training task has been accepted and started in the background."},
        400: {"description": "The 'texts' and 'labels' lists have mismatched lengths."},
        422: {"description": "Validation error, e.g., the request body is malformed or lists are empty."},
    },
)
async def train_from_data(
    request: TrainDirectRequest,
    background_tasks: BackgroundTasks,
    training_service: TrainingService = Depends(get_training_service),
):
    if len(request.texts) != len(request.labels):
        raise HTTPException(
            status_code=400, detail="The 'texts' and 'labels' lists must have the same number of items.")

    background_tasks.add_task(
        _run_training_task, training_service, request.texts, request.labels)

    return TrainingResponse(message="Training has been accepted and started in the background. Check logs for completion status.")