"""Main classification API for domain routing."""

import logging
from typing import Optional
import time

import torch

from .model_loader import get_model, get_tokenizer
from .schemas import ClassificationResult, PredictionDetail

logger = logging.getLogger(__name__)


def classify_domain(
    text: str, threshold: float = 0.7, multi_label: bool = False, log_latency: bool = False
) -> ClassificationResult:
    """Classify text into domain(s) for chatbot routing.

    This is the main API for the denes-router-classifier package. It uses a fine-tuned
    BERT multilingual model to classify user queries into one of three domains:
    - Web_Search: Explicit search requests (verbs: busca, encuentra, investiga)
    - TI: Technical support queries (hardware, software, network issues)
    - Generic: General knowledge questions answerable by LLM (no web search needed)

    Args:
        text: Input text to classify (user query)
        threshold: Confidence threshold for predictions (default: 0.7)
            - Predictions below this threshold trigger includes_generic=True
        multi_label: Return all domains above threshold (default: False)
            - False: Returns only the highest confidence domain
            - True: Returns all domains exceeding the threshold
        log_latency: Log inference latency to logger (default: False)

    Returns:
        ClassificationResult with:
            - primary: Domain with highest confidence
            - confidence: Confidence score (0-1) for primary domain
            - includes_generic: True if confidence < threshold (consider Generic fallback)
            - all_predictions: List of all predictions if multi_label=True

    Examples:
        >>> # Single-label classification (default)
        >>> result = classify_domain("Busca el clima en Asunción")
        >>> result.primary
        'Web_Search'
        >>> result.confidence
        0.9984

        >>> # Multi-label with lower threshold
        >>> result = classify_domain(
        ...     "Busca ayuda para mi PC",
        ...     threshold=0.5,
        ...     multi_label=True
        ... )
        >>> len(result.all_predictions)
        2  # Both Web_Search and TI above 0.5

        >>> # Low confidence triggers Generic fallback
        >>> result = classify_domain("Hmm...", threshold=0.7)
        >>> result.includes_generic
        True

    Raises:
        ValueError: If text is empty or threshold is not in [0, 1]
        RuntimeError: If model loading fails
    """
    # Input validation
    if not text or not text.strip():
        raise ValueError("Input text cannot be empty")

    if not 0 <= threshold <= 1:
        raise ValueError(f"Threshold must be between 0 and 1, got {threshold}")

    # Start latency measurement
    start_time = time.perf_counter() if log_latency else None

    # Load model and tokenizer (singleton, cached after first call)
    tokenizer = get_tokenizer()
    model = get_model()

    # Tokenize input
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=160, padding=False)

    # Inference (no gradient computation)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=1)[0]

    # Get all predictions above threshold
    predictions: list[PredictionDetail] = []
    for idx, prob in enumerate(probs):
        prob_value = prob.item()
        if prob_value >= threshold:
            domain_name = model.config.id2label[idx]
            predictions.append(PredictionDetail(domain=domain_name, confidence=prob_value))

    # Sort by confidence (descending)
    predictions.sort(key=lambda x: x.confidence, reverse=True)

    # Determine max confidence and includes_generic flag
    max_confidence = probs.max().item()
    includes_generic = max_confidence < threshold

    # Log latency if requested
    if log_latency and start_time is not None:
        latency_ms = (time.perf_counter() - start_time) * 1000
        logger.info(f"Classification latency: {latency_ms:.2f}ms")

    # Build result based on mode
    if multi_label:
        # Multi-label mode: return all predictions
        primary_domain = predictions[0].domain if predictions else "Generic"
        primary_confidence = predictions[0].confidence if predictions else max_confidence

        return ClassificationResult(
            primary=primary_domain,
            confidence=primary_confidence,
            includes_generic=includes_generic,
            all_predictions=predictions if predictions else None,
        )
    else:
        # Single-label mode: return only primary domain
        if predictions:
            return ClassificationResult(
                primary=predictions[0].domain,
                confidence=predictions[0].confidence,
                includes_generic=includes_generic,
            )
        else:
            # No predictions above threshold, fallback to Generic
            return ClassificationResult(primary="Generic", confidence=max_confidence, includes_generic=True)


def classify_batch(
    texts: list[str], threshold: float = 0.7, multi_label: bool = False, batch_size: int = 32
) -> list[ClassificationResult]:
    """Classify multiple texts in batch for better throughput.

    More efficient than calling classify_domain() in a loop for multiple texts.

    Args:
        texts: List of input texts to classify
        threshold: Confidence threshold for predictions (default: 0.7)
        multi_label: Return all domains above threshold (default: False)
        batch_size: Batch size for inference (default: 32)

    Returns:
        List of ClassificationResult objects (same order as input)

    Examples:
        >>> texts = [
        ...     "Busca el clima",
        ...     "No tengo internet",
        ...     "¿Cuál es la capital de Francia?"
        ... ]
        >>> results = classify_batch(texts)
        >>> [r.primary for r in results]
        ['Web_Search', 'TI', 'Generic']

    Raises:
        ValueError: If texts is empty or contains invalid entries
    """
    if not texts:
        raise ValueError("Input texts list cannot be empty")

    if any(not text or not text.strip() for text in texts):
        raise ValueError("Input texts cannot contain empty strings")

    # Load model and tokenizer
    tokenizer = get_tokenizer()
    model = get_model()

    results: list[ClassificationResult] = []

    # Process in batches
    for i in range(0, len(texts), batch_size):
        batch_texts = texts[i : i + batch_size]

        # Tokenize batch
        inputs = tokenizer(batch_texts, return_tensors="pt", truncation=True, max_length=160, padding=True)

        # Inference
        with torch.no_grad():
            outputs = model(**inputs)
            probs_batch = torch.softmax(outputs.logits, dim=1)

        # Process each item in batch
        for probs in probs_batch:
            # Get predictions above threshold
            predictions: list[PredictionDetail] = []
            for idx, prob in enumerate(probs):
                prob_value = prob.item()
                if prob_value >= threshold:
                    domain_name = model.config.id2label[idx]
                    predictions.append(PredictionDetail(domain=domain_name, confidence=prob_value))

            predictions.sort(key=lambda x: x.confidence, reverse=True)

            max_confidence = probs.max().item()
            includes_generic = max_confidence < threshold

            if multi_label:
                primary_domain = predictions[0].domain if predictions else "Generic"
                primary_confidence = predictions[0].confidence if predictions else max_confidence

                results.append(
                    ClassificationResult(
                        primary=primary_domain,
                        confidence=primary_confidence,
                        includes_generic=includes_generic,
                        all_predictions=predictions if predictions else None,
                    )
                )
            else:
                if predictions:
                    results.append(
                        ClassificationResult(
                            primary=predictions[0].domain,
                            confidence=predictions[0].confidence,
                            includes_generic=includes_generic,
                        )
                    )
                else:
                    results.append(
                        ClassificationResult(primary="Generic", confidence=max_confidence, includes_generic=True)
                    )

    return results
