"""
Workflow orchestration for RedenLab ML SDK.

Provides high-level functions that compose session + client for common patterns.
This is where upload + inference + aggregation logic lives, keeping the client
class focused on single-asset operations.
"""

import statistics
from collections.abc import Callable
from typing import Any

from .exceptions import InferenceError

# Type alias for aggregator function
Aggregator = Callable[[list[dict[str, Any]]], dict[str, Any]]


def run_chunked_inference(
    session: Any,  # RedenLabSession
    client: Any,  # InferenceClient
    file_path: str,
    chunk_duration: float = 50.0,
    progress_callback: Callable[[dict[str, Any]], None] | None = None,
) -> list[dict[str, Any]]:
    """
    Upload with chunking, run inference on each chunk, return raw results.
    The caller decides how to aggregate (or not aggregate) the results.

    Args:
        session: RedenLabSession instance
        client: InferenceClient instance
        file_path: Path to the audio file
        chunk_duration: Duration of each chunk in seconds (default: 50)
        progress_callback: Optional callback for progress updates

    Returns:
        List of result dictionaries, each with:
        - All original fields from inference result
        - _chunk_index: Index of this chunk (0-based)
        - _total_chunks: Total number of chunks
        - _failed_chunks: List of failed chunk indices (on last result only)

    Raises:
        InferenceError: If all chunks fail
        TimeoutError: If inference times out
        APIError: If API communication fails

    Example:
        >>> # Get raw chunk results
        >>> results = run_chunked_inference(session, client, "long_audio.wav")
        >>> for r in results:
        ...     print(f"Chunk {r['_chunk_index']}: {r['result']}")
        >>>
        >>> # Then aggregate however you want
        >>> aggregated = aggregate_results(results, "intelligibility_score")
    """
    # Upload with chunking (backend tracks relationships)
    chunks = session.upload_chunked(file_path, chunk_duration)
    total_chunks = len(chunks)

    # Single chunk - still return as list for consistent interface
    if len(chunks) == 1:
        result = client.predict(chunks[0], progress_callback=progress_callback)
        result["_chunk_index"] = 0
        result["_total_chunks"] = 1
        return [result]

    # Multiple chunks - submit all, poll all
    job_ids = [client.submit(chunk) for chunk in chunks]

    results = []
    failed_chunks = []

    for i, job_id in enumerate(job_ids):
        try:
            result = client.poll(job_id, progress_callback=progress_callback)
            if result["status"] == "completed":
                # Inject chunk metadata
                result["_chunk_index"] = i
                result["_total_chunks"] = total_chunks
                results.append(result)
            else:
                failed_chunks.append(i)
        except Exception:
            # Track failed chunk but continue polling others
            failed_chunks.append(i)

    if not results:
        raise InferenceError(
            f"All {total_chunks} chunks failed. Failed chunk indices: {failed_chunks}"
        )

    # Attach failed_chunks info to last result for visibility
    if failed_chunks:
        results[-1]["_failed_chunks"] = failed_chunks

    return results


def aggregate_results(
    results: list[dict[str, Any]],
    score_key: str,
    reduce_fn: Callable[[list[float]], float] | None = None,
) -> dict[str, Any]:
    """
    Aggregate scores from multiple chunk results.

    Args:
        results: List of successful job results (from run_chunked_inference)
        score_key: Key to extract score from results (e.g., 'intelligibility_score')
        reduce_fn: Function to reduce scores to single value.
                  Signature: (list[float]) -> float
                  Default: mean (sum/len)

    Returns:
        Aggregated result dictionary with:
        - job_id: Prefixed with "chunked_"
        - status: "completed" or "partial" (if some chunks failed)
        - result: Contains aggregated score, std_dev, chunk_scores, metadata
        - created_at, completed_at: Timestamps from chunks

    Example:
        >>> results = run_chunked_inference(session, client, "long.wav")
        >>>
        >>> # Default aggregation (mean)
        >>> agg = aggregate_results(results, "intelligibility_score")
        >>>
        >>> # Custom aggregation (median)
        >>> import statistics
        >>> agg = aggregate_results(results, "intelligibility_score", reduce_fn=statistics.median)
        >>>
        >>> # Custom aggregation (max)
        >>> agg = aggregate_results(results, "intelligibility_score", reduce_fn=max)
    """
    if not results:
        raise ValueError("No results to aggregate")

    # Default reduce function is mean
    if reduce_fn is None:

        def reduce_fn(scores):
            return sum(scores) / len(scores)

        aggregation_method = "mean"
    else:
        aggregation_method = "custom"

    # Extract scores
    scores = [r["result"][score_key] for r in results]

    # Apply reduce function
    aggregated_score = reduce_fn(scores)

    # Always compute std_dev for observability (even with custom reduce_fn)
    std_dev = statistics.stdev(scores) if len(scores) > 1 else 0.0

    # Check for failed chunks
    failed_chunks = results[-1].get("_failed_chunks", [])
    total_chunks = results[0].get("_total_chunks", len(results))

    # Determine status
    status = "completed" if not failed_chunks else "partial"

    # Build aggregated result
    aggregated = {
        "job_id": f"chunked_{results[0]['job_id']}",
        "status": status,
        "result": {
            score_key: aggregated_score,
            "std_dev": std_dev,
            "chunk_scores": scores,
            "num_chunks": len(results),
            "total_chunks": total_chunks,
            "aggregation_method": aggregation_method,
        },
        "created_at": min(r["created_at"] for r in results),
        "completed_at": max(r.get("completed_at", r["created_at"]) for r in results),
    }

    # Add failure info if applicable
    if failed_chunks:
        aggregated["result"]["failed_chunks"] = failed_chunks
        aggregated["result"]["success_rate"] = f"{len(results)}/{total_chunks}"

    return aggregated


def _default_aggregator_for(client: Any) -> Aggregator | None:
    """
    Build a default aggregator for the given client.

    Returns None if client doesn't have a score_key (e.g., TranscribeClient).
    This signals that auto-chunking should not be triggered.

    Args:
        client: InferenceClient instance

    Returns:
        Aggregator function or None
    """
    score_key = getattr(client, "score_key", None)
    if score_key is None:
        return None

    # Return an aggregator that uses this client's score_key with default mean
    def aggregator(results: list[dict[str, Any]]) -> dict[str, Any]:
        return aggregate_results(results, score_key)

    return aggregator


def run_inference(
    session: Any,  # RedenLabSession
    client: Any,  # InferenceClient
    file_path: str,
    chunk_duration: float | None = None,
    aggregator: Aggregator | None = None,
    **model_params: Any,
) -> dict[str, Any]:
    """
    Upload file and run inference. Handles chunking transparently.

    This is the primary entry point for most users (Level 1).
    It automatically chunks long files based on client.chunk_threshold,
    runs inference, and aggregates results using the provided aggregator.

    Args:
        session: RedenLabSession instance
        client: InferenceClient instance
        file_path: Path to the audio file
        chunk_duration: Duration of each chunk in seconds (optional).
                       Only used if file exceeds client.chunk_threshold.
                       Defaults to 50s if not specified.
        aggregator: Function to aggregate chunk results.
                   Signature: (list[dict]) -> dict
                   If None, uses _default_aggregator_for(client).
                   If client has no score_key, chunking is not auto-triggered.
        **model_params: Model-specific parameters passed to predict()

    Returns:
        Inference result dictionary (aggregated if chunked)

    Raises:
        ValidationError: If file is invalid
        InferenceError: If inference fails
        TimeoutError: If inference times out
        APIError: If API communication fails

    Example:
        >>> from redenlab_extract import RedenLabSession, intelligibility_client, run_inference
        >>> session = RedenLabSession(api_key="...")
        >>> client = intelligibility_client(session, model_name="als-intelligibility-v1")
        >>>
        >>> # Short file - uploads and predicts directly
        >>> result = run_inference(session, client, "short.wav")
        >>>
        >>> # Long file - chunks, predicts each, aggregates (all transparent!)
        >>> result = run_inference(session, client, "long_recording.wav")
        >>> print(result['result']['intelligibility_score'])  # Aggregated score
        >>>
        >>> # Custom aggregator (e.g., median instead of mean)
        >>> import statistics
        >>> def median_aggregator(results):
        ...     return aggregate_results(results, "intelligibility_score", reduce_fn=statistics.median)
        >>> result = run_inference(session, client, "long.wav", aggregator=median_aggregator)
    """
    # Check duration to decide if chunking is needed
    try:
        duration = session._get_audio_duration(file_path)
    except Exception:
        duration = None

    # Determine effective aggregator
    effective_aggregator = aggregator or _default_aggregator_for(client)

    # Auto-chunk only if:
    # 1. Duration exceeds threshold, AND
    # 2. We have an aggregator (either provided or from client.score_key)
    if duration and duration > client.chunk_threshold and effective_aggregator:
        # Get raw chunk results
        results = run_chunked_inference(
            session=session,
            client=client,
            file_path=file_path,
            chunk_duration=chunk_duration or 50.0,
        )
        # Aggregate using the effective aggregator
        return effective_aggregator(results)

    # Short file, no aggregator, or duration unavailable - use simple workflow
    asset = session.upload(file_path)
    result: dict[str, Any] = client.predict(asset, **model_params)
    return result
