"""
Synchronous Composo client

This module provides a synchronous Composo client for evaluating chat message quality.
Composo is an AI evaluation platform that can analyze chat conversations and tool calls,
and provide scores and explanations based on specified criteria.

Main features:
- Evaluate chat message quality and relevance
- Support multiple evaluation criteria (single or multiple criteria)
- Handle tool calls and system messages
- Provide retry mechanisms and error handling
"""

from typing import List, Dict, Any, Optional, Union

from ..models.api_types import (
    MultiAgentTrace,
    MultiAgentTraceResponse,
)
from ..models.common_types import ModelCore
from .base import BaseClient, REWARD_ENDPOINT, TRACE_ENDPOINT
from .types import MessagesType, ToolsType, ResultType
from ..models import EvaluationResponse
from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    wait_random,
    retry_if_exception_type,
)


class Composo(BaseClient):
    """Synchronous Composo client.

    This client is used for synchronous calls to the Composo API for message evaluation.
    Suitable for single evaluations or small batch evaluation scenarios, providing a simple and easy-to-use interface.

    Key features:
        - Synchronous API calls, easy to integrate into existing code
        - Automatic retry mechanism, improving request success rate
        - Support for context managers, automatic resource management
        - Support for multiple evaluation criteria

    Inherits from BaseClient, which provides:
        - api_key: Composo API key for authentication
        - base_url: API base URL, defaults to official platform address
        - num_retries: Number of retries on request failure, defaults to 1
        - model_core: Optional model core identifier for specifying evaluation model
        - timeout: Request timeout in seconds, defaults to 60 seconds
    """

    def __enter__(self):
        """Context manager entry"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - close HTTP client"""
        self.close()

    def _make_request(
        self, request_data: Dict[str, Any], endpoint: str
    ) -> Dict[str, Any]:
        """Make HTTP request with retry logic using tenacity (exponential backoff base 2 + jitter)"""

        @retry(
            stop=stop_after_attempt(self.num_retries),
            wait=wait_exponential(multiplier=1, min=1, max=60) + wait_random(0, 1),
            retry=retry_if_exception_type(Exception),
            reraise=True,
        )
        def do_request():
            return self._post(
                endpoint=endpoint,
                data=request_data,
                headers=self._build_headers(),
            )

        return do_request()

    def _evaluate_single_criterion(
        self,
        messages: Optional[MessagesType],
        criterion: Optional[str],
        system: Optional[str],
        tools: ToolsType,
        result: ResultType,
        block: bool = True,
        tags: Optional[Dict[str, str]] = None,
        evaluate_latest: Optional[bool] = None,
        evaluator_name: Optional[str] = None,
        input_: Optional[Any] = None,
    ) -> EvaluationResponse:
        """Evaluate a single criterion synchronously"""
        evaluation_request = self._prepare_evaluation_request(
            messages,
            system,
            tools,
            result,
            criterion,
            tags,
            evaluate_latest,
            evaluator_name=evaluator_name,
            input_=input_,
        )
        request_data = evaluation_request.model_dump(exclude_none=True)

        # Determine endpoint and make request
        endpoint = self._get_endpoint(evaluation_request, criterion)
        if not block:
            endpoint = f"{endpoint}/background"
        response_data = self._make_request(request_data, endpoint)

        # If non-blocking, return response as-is (contains task_id)
        if not block:
            return response_data

        # Process response using shared logic
        return self._process_evaluation_response(response_data, criterion)

    def evaluate(
        self,
        messages: Optional[MessagesType] = None,
        input: Optional[Any] = None,
        system: Optional[str] = None,
        tools: ToolsType = None,
        result: ResultType = None,
        criteria: Optional[Union[str, List[str]]] = None,
        block: bool = True,
        tags: Optional[Dict[str, str]] = None,
        evaluate_latest: Optional[bool] = None,
        evaluator_name: Optional[Union[str, List[str]]] = None,
    ) -> Union[EvaluationResponse, List[EvaluationResponse]]:
        """Evaluate messages with optional criteria.

        Args:
            messages: List of chat messages to evaluate.
                Format: [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hello!"}]
                Mutually exclusive with input.
            input: Raw input passed to the OpenAI Responses API (str or list).
                Auto-converted to messages format. Mutually exclusive with messages.
            criteria: Evaluation criteria text (str or list of str). Mutually exclusive with evaluator_name.
            system: Optional system message to set AI behavior and context.
            tools: Optional list of tool definitions for evaluating tool calls.
            result: Optional LLM result to append to the conversation.
            block: If False, returns dict with task_id instead of blocking for results.
            tags: Optional key-value pairs to tag and categorize the request.
            evaluate_latest: Whether to evaluate only the latest response or all responses.
                If not provided, API defaults apply (True for chat evaluations).
            evaluator_name: Name of an existing evaluator (str or list of str).
                Mutually exclusive with criteria.

        Returns:
            EvaluationResponse or list of EvaluationResponse objects.
            If block=False, returns dict(s) with task_id.
        """

        if messages is not None and input is not None:
            raise ValueError("Provide either 'messages' or 'input', not both")
        if messages is None and input is None:
            raise ValueError("Either 'messages' or 'input' is required")

        # Normalize to lists
        if isinstance(criteria, str):
            criteria = [criteria]
        if isinstance(evaluator_name, str):
            evaluator_name = [evaluator_name]

        if criteria is not None and evaluator_name is not None:
            raise ValueError("Provide either 'criteria' or 'evaluator_name', not both")
        if criteria is None and evaluator_name is None:
            raise ValueError("Either 'criteria' or 'evaluator_name' is required")

        results = self._evaluate_multiple_criteria(
            messages,
            criteria,
            system,
            tools,
            result,
            block,
            tags,
            evaluate_latest,
            evaluator_names=evaluator_name,
            input_=input,
        )

        n = len(criteria or evaluator_name or [None])
        return results[0] if n == 1 else results

    def _evaluate_multiple_criteria(
        self,
        messages: Optional[MessagesType],
        criteria: Optional[List[str]],
        system: Optional[str],
        tools: ToolsType,
        result: ResultType,
        block: bool = True,
        tags: Optional[Dict[str, str]] = None,
        evaluate_latest: Optional[bool] = None,
        evaluator_names: Optional[List[str]] = None,
        input_: Optional[Any] = None,
    ) -> List[EvaluationResponse]:
        """Evaluate multiple criteria sequentially"""
        if criteria:
            pairs = [(c, None) for c in criteria]
        else:
            pairs = [(None, n) for n in evaluator_names]

        results = []
        for ec, en in pairs:
            results.append(
                self._evaluate_single_criterion(
                    messages=messages,
                    criterion=ec,
                    system=system,
                    tools=tools,
                    result=result,
                    block=block,
                    tags=tags,
                    evaluate_latest=evaluate_latest,
                    evaluator_name=en,
                    input_=input_,
                )
            )
        return results

    def _evaluate_trace(
        self,
        trace: MultiAgentTrace,
        model_core: Optional[ModelCore],
        criteria: Optional[List[str]],
        block: bool = True,
        tags: Optional[Dict[str, str]] = None,
        evaluate_latest: Optional[bool] = None,
        evaluator_names: Optional[List[str]] = None,
    ) -> List[MultiAgentTraceResponse]:
        if criteria:
            pairs = [(c, None) for c in criteria]
        else:
            pairs = [(None, n) for n in evaluator_names]

        res = []
        for ec, en in pairs:
            request_info = self._prepare_trace_request(
                trace, ec, model_core, tags, evaluate_latest, evaluator_name=en
            )
            endpoint = request_info["endpoint"]
            if not block:
                endpoint = f"{endpoint}/background"
            response_data = self._make_request(request_info["request_data"], endpoint)
            if not block:
                res.append(response_data)
            else:
                trace_response = self._process_trace_response(response_data)
                res.append(trace_response)
        return res

    def evaluate_trace(
        self,
        trace: MultiAgentTrace,
        criteria: Optional[Union[str, List[str]]] = None,
        model_core: Optional[ModelCore] = None,
        block: bool = True,
        tags: Optional[Dict[str, str]] = None,
        evaluate_latest: Optional[bool] = None,
        evaluator_name: Optional[Union[str, List[str]]] = None,
    ) -> Union[MultiAgentTraceResponse, List[MultiAgentTraceResponse]]:
        """Evaluate trace with optional criteria.

        Args:
            trace: MultiAgentTrace object to evaluate.
            criteria: Evaluation criteria text (str or list of str). Mutually exclusive with evaluator_name.
            model_core: Optional model core identifier for specifying evaluation model.
            block: If False, returns dict(s) with task_id instead of blocking for results.
            tags: Optional key-value pairs to tag and categorize the request.
            evaluate_latest: Whether to evaluate only the latest response or all responses.
                If not provided, API defaults apply (False for trace evaluations).
            evaluator_name: Name of an existing evaluator (str or list of str).
                Mutually exclusive with criteria.

        Returns:
            MultiAgentTraceResponse or list of MultiAgentTraceResponse objects.
            If block=False, returns dict(s) with task_id.
        """

        # Normalize to lists
        if isinstance(criteria, str):
            criteria = [criteria]
        if isinstance(evaluator_name, str):
            evaluator_name = [evaluator_name]

        if criteria is not None and evaluator_name is not None:
            raise ValueError("Provide either 'criteria' or 'evaluator_name', not both")
        if criteria is None and evaluator_name is None:
            raise ValueError("Either 'criteria' or 'evaluator_name' is required")

        # Use model_core from parameter, or self.model_core, or None (will use TraceRewardRequest default)
        final_model_core = model_core if model_core is not None else self.model_core

        results = self._evaluate_trace(
            trace,
            final_model_core,
            criteria,
            block,
            tags,
            evaluate_latest,
            evaluator_names=evaluator_name,
        )

        n = len(criteria or evaluator_name or [None])
        return results[0] if n == 1 else results
