import sys
import json
from typing import Any, Dict, List, Union, Literal, TypeVar, Callable, Optional, cast
from functools import wraps

import click
from tabulate import tabulate

from together import APIError, Together, TogetherError
from together._types import omit
from together.lib.utils.serializer import datetime_serializer
from together.types.eval_create_params import (
    ParametersEvaluationScoreParameters,
    ParametersEvaluationCompareParameters,
    ParametersEvaluationClassifyParameters,
    ParametersEvaluationScoreParametersJudge,
    ParametersEvaluationCompareParametersJudge,
    ParametersEvaluationClassifyParametersJudge,
    ParametersEvaluationScoreParametersModelToEvaluate,
    ParametersEvaluationClassifyParametersModelToEvaluate,
    ParametersEvaluationCompareParametersModelAEvaluationModelRequest,
    ParametersEvaluationCompareParametersModelBEvaluationModelRequest,
)


def print_api_error(e: Union[APIError, TogetherError]) -> None:
    if isinstance(e, APIError):
        error_details = cast(Dict[str, Any], e.body)["error"]["message"]

        if error_details and ("credentials" in error_details.lower() or "authentication" in error_details.lower()):
            click.echo("Error: Invalid API key or authentication failed", err=True)
        else:
            click.echo(f"Error: {error_details}", err=True)

    click.echo(f"Error: {e}", err=True)
    return


F = TypeVar("F", bound=Callable[..., Any])


def handle_api_errors(f: F) -> F:
    """Decorator to handle common API errors in CLI commands."""

    @wraps(f)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        try:
            return f(*args, **kwargs)
        except APIError as e:
            print_api_error(e)
            sys.exit(1)
        except TogetherError as e:
            print_api_error(e)
            sys.exit(1)
        except Exception as e:
            click.echo(f"Error: An unexpected error occurred - {str(e)}", err=True)
            sys.exit(1)

    return wrapper  # type: ignore


@click.group()
@click.pass_context
def evals(ctx: click.Context) -> None:
    """Evals API commands"""
    pass


@evals.command()
@click.pass_context
@click.option(
    "--type",
    type=click.Choice(["classify", "score", "compare"]),
    required=True,
    help="Type of evaluation to create.",
)
@click.option(
    "--judge-model",
    type=str,
    required=True,
    help="Name or URL of the judge model to use for evaluation.",
)
@click.option(
    "--judge-model-source",
    type=click.Choice(["serverless", "dedicated", "external"]),
    required=True,
    help="Source of the judge model.",
)
@click.option(
    "--judge-external-api-token",
    type=str,
    required=False,
    help="Optional external API token for the judge model.",
)
@click.option(
    "--judge-external-base-url",
    type=str,
    required=False,
    help="Optional external base URLs for the judge model.",
)
@click.option(
    "--judge-system-template",
    type=str,
    required=True,
    help="System template for the judge model.",
)
@click.option(
    "--input-data-file-path",
    type=str,
    required=True,
    help="Path to the input data file.",
)
@click.option(
    "--model-field",
    type=str,
    help="Name of the field in the input file contaning text generated by the model."
    "Can not be used when model-a-name and other model config parameters are specified",
)
@click.option(
    "--model-to-evaluate",
    type=str,
    help="Model name when using the detailed config",
)
@click.option(
    "--model-to-evaluate-source",
    type=click.Choice(["serverless", "dedicated", "external"]),
    help="Source of the model to evaluate.",
)
@click.option(
    "--model-to-evaluate-external-api-token",
    type=str,
    help="Optional external API token for the model to evaluate.",
)
@click.option(
    "--model-to-evaluate-external-base-url",
    type=str,
    help="Optional external base URL for the model to evaluate.",
)
@click.option(
    "--model-to-evaluate-max-tokens",
    type=int,
    help="Max tokens for model-to-evaluate",
)
@click.option(
    "--model-to-evaluate-temperature",
    type=float,
    help="Temperature for model-to-evaluate",
)
@click.option(
    "--model-to-evaluate-system-template",
    type=str,
    help="System template for model-to-evaluate",
)
@click.option(
    "--model-to-evaluate-input-template",
    type=str,
    help="Input template for model-to-evaluate",
)
@click.option(
    "--labels",
    type=str,
    help="Classification labels - comma-separated list",
)
@click.option(
    "--pass-labels",
    type=str,
    help="Labels considered as passing (required for classify type). A comma-separated list.",
)
@click.option(
    "--min-score",
    type=float,
    help="Minimum score value (required for score type).",
)
@click.option(
    "--max-score",
    type=float,
    help="Maximum score value (required for score type).",
)
@click.option(
    "--pass-threshold",
    type=float,
    help="Threshold score for passing (required for score type).",
)
@click.option(
    "--model-a-field",
    type=str,
    help="Name of the field in the input file containing text generated by Model A. \
        Can not be used when model-a-name and other model config parameters are specified",
)
@click.option(
    "--model-a",
    type=str,
    help="Model name or URL for model A when using detailed config.",
)
@click.option(
    "--model-a-source",
    type=click.Choice(["serverless", "dedicated", "external"]),
    help="Source of model A.",
)
@click.option(
    "--model-a-external-api-token",
    type=str,
    help="Optional external API token for model A.",
)
@click.option(
    "--model-a-external-base-url",
    type=str,
    help="Optional external base URL for model A.",
)
@click.option(
    "--model-a-max-tokens",
    type=int,
    help="Max tokens for model A.",
)
@click.option(
    "--model-a-temperature",
    type=float,
    help="Temperature for model A.",
)
@click.option(
    "--model-a-system-template",
    type=str,
    help="System template for model A.",
)
@click.option(
    "--model-a-input-template",
    type=str,
    help="Input template for model A.",
)
@click.option(
    "--model-b-field",
    type=str,
    help="Name of the field in the input file containing text generated by Model B.\
          Can not be used when model-b-name and other model config parameters are specified",
)
@click.option(
    "--model-b",
    type=str,
    help="Model name or URL for model B when using detailed config.",
)
@click.option(
    "--model-b-source",
    type=click.Choice(["serverless", "dedicated", "external"]),
    help="Source of model B.",
)
@click.option(
    "--model-b-external-api-token",
    type=str,
    help="Optional external API token for model B.",
)
@click.option(
    "--model-b-external-base-url",
    type=str,
    help="Optional external base URL for model B.",
)
@click.option(
    "--model-b-max-tokens",
    type=int,
    help="Max tokens for model B.",
)
@click.option(
    "--model-b-temperature",
    type=float,
    help="Temperature for model B.",
)
@click.option(
    "--model-b-system-template",
    type=str,
    help="System template for model B.",
)
@click.option(
    "--model-b-input-template",
    type=str,
    help="Input template for model B.",
)
@handle_api_errors
def create(
    ctx: click.Context,
    type: Literal["classify", "score", "compare"],
    judge_model: str,
    judge_model_source: Literal["serverless", "dedicated", "external"],
    judge_system_template: str,
    judge_external_api_token: Optional[str],
    judge_external_base_url: Optional[str],
    input_data_file_path: str,
    model_field: Optional[str],
    model_to_evaluate: Optional[str],
    model_to_evaluate_source: Optional[str],
    model_to_evaluate_external_api_token: Optional[str],
    model_to_evaluate_external_base_url: Optional[str],
    model_to_evaluate_max_tokens: Optional[int],
    model_to_evaluate_temperature: Optional[float],
    model_to_evaluate_system_template: Optional[str],
    model_to_evaluate_input_template: Optional[str],
    labels: str,
    pass_labels: str,
    min_score: Optional[float],
    max_score: Optional[float],
    pass_threshold: Optional[float],
    model_a_field: Optional[str],
    model_a: Optional[str],
    model_a_source: Optional[str],
    model_a_external_api_token: Optional[str],
    model_a_external_base_url: Optional[str],
    model_a_max_tokens: Optional[int],
    model_a_temperature: Optional[float],
    model_a_system_template: Optional[str],
    model_a_input_template: Optional[str],
    model_b_field: Optional[str],
    model_b: Optional[str],
    model_b_source: Optional[str],
    model_b_external_api_token: Optional[str],
    model_b_external_base_url: Optional[str],
    model_b_max_tokens: Optional[int],
    model_b_temperature: Optional[float],
    model_b_system_template: Optional[str],
    model_b_input_template: Optional[str],
) -> None:
    """Create a new evaluation job"""

    client: Together = ctx.obj

    # Convert strings to lists for labels
    labels_list = labels.split(",") if labels else None
    pass_labels_list = pass_labels.split(",") if pass_labels else None

    # Build model configurations
    model_to_evaluate_final: Union[Dict[str, Any], None, str] = None

    # Check if any config parameters are provided
    config_params_provided = any(
        [
            model_to_evaluate,
            model_to_evaluate_source,
            model_to_evaluate_max_tokens,
            model_to_evaluate_temperature,
            model_to_evaluate_system_template,
            model_to_evaluate_input_template,
        ]
    )

    if model_field:
        # Simple mode: model_field is provided
        if config_params_provided:
            raise click.BadParameter(
                "Cannot specify both --model-field and --model-to-evaluate-* parameters. "
                "Use either --model-field alone if your input file has pre-generated responses, "
                "or config parameters if you want to generate it on our end"
            )
        model_to_evaluate_final = model_field
    elif config_params_provided:
        # Config mode: config parameters are provided
        model_to_evaluate_final = {
            "model": model_to_evaluate,
            "model_source": model_to_evaluate_source,
            "max_tokens": model_to_evaluate_max_tokens,
            "temperature": model_to_evaluate_temperature,
            "system_template": model_to_evaluate_system_template,
            "input_template": model_to_evaluate_input_template,
        }
        if model_to_evaluate_external_api_token:
            model_to_evaluate_final["external_api_token"] = model_to_evaluate_external_api_token
        if model_to_evaluate_external_base_url:
            model_to_evaluate_final["external_base_url"] = model_to_evaluate_external_base_url

    # Build model-a configuration
    model_a_final: Union[Dict[str, Any], None, str] = None
    model_a_config_params = [
        model_a,
        model_a_source,
        model_a_max_tokens,
        model_a_temperature,
        model_a_system_template,
        model_a_input_template,
    ]

    if model_a_field is not None:
        # Simple mode: model_a_field is provided
        if any(model_a_config_params):
            raise click.BadParameter(
                "Cannot specify both --model-a-field and config parameters (--model-a-name, etc.). "
                "Use either --model-a-field alone if your input file has pre-generated responses, "
                "or config parameters if you want to generate it on our end"
            )
        model_a_final = model_a_field
    elif any(model_a_config_params):
        # Config mode: config parameters are provided
        model_a_final = {
            "model": model_a,
            "model_source": model_a_source,
            "max_tokens": model_a_max_tokens,
            "temperature": model_a_temperature,
            "system_template": model_a_system_template,
            "input_template": model_a_input_template,
        }
        if model_a_external_api_token:
            model_a_final["external_api_token"] = model_a_external_api_token
        if model_a_external_base_url:
            model_a_final["external_base_url"] = model_a_external_base_url

    # Build model-b configuration
    model_b_final: Union[Dict[str, Any], None, str] = None
    model_b_config_params = [
        model_b,
        model_b_source,
        model_b_max_tokens,
        model_b_temperature,
        model_b_system_template,
        model_b_input_template,
    ]

    if model_b_field is not None:
        # Simple mode: model_b_field is provided
        if any(model_b_config_params):
            raise click.BadParameter(
                "Cannot specify both --model-b-field and config parameters (--model-b-name, etc.). "
                "Use either --model-b-field alone if your input file has pre-generated responses, "
                "or config parameters if you want to generate it on our end"
            )
        model_b_final = model_b_field
    elif any(model_b_config_params):
        # Config mode: config parameters are provided
        model_b_final = {
            "model": model_b,
            "model_source": model_b_source,
            "max_tokens": model_b_max_tokens,
            "temperature": model_b_temperature,
            "system_template": model_b_system_template,
            "input_template": model_b_input_template,
        }
        if model_b_external_api_token:
            model_b_final["external_api_token"] = model_b_external_api_token
        if model_b_external_base_url:
            model_b_final["external_base_url"] = model_b_external_base_url

    judge_config = _build_judge(
        type, judge_model, judge_model_source, judge_system_template, judge_external_api_token, judge_external_base_url
    )

    if type == "classify":
        response = client.evals.create(
            type=type,
            parameters=ParametersEvaluationClassifyParameters(
                input_data_file_path=input_data_file_path,
                judge=judge_config,
                labels=labels_list or [],
                pass_labels=pass_labels_list or [],
                model_to_evaluate=cast(ParametersEvaluationClassifyParametersModelToEvaluate, model_to_evaluate_final),
            ),
        )
    elif type == "score":
        if max_score is None or min_score is None or pass_threshold is None:
            raise TogetherError("max_score, min_score, and pass_threshold are required for score type")

        response = client.evals.create(
            type="score",
            parameters=ParametersEvaluationScoreParameters(
                input_data_file_path=input_data_file_path,
                judge=judge_config,
                max_score=max_score,
                min_score=min_score,
                pass_threshold=pass_threshold,
                model_to_evaluate=cast(ParametersEvaluationScoreParametersModelToEvaluate, model_to_evaluate_final),
            ),
        )
    elif type == "compare":
        response = client.evals.create(
            type=type,
            parameters=ParametersEvaluationCompareParameters(
                input_data_file_path=input_data_file_path,
                judge=judge_config,
                model_a=cast(ParametersEvaluationCompareParametersModelAEvaluationModelRequest, model_a_final),
                model_b=cast(ParametersEvaluationCompareParametersModelBEvaluationModelRequest, model_b_final),
            ),
        )

    click.echo(json.dumps(response.model_dump(exclude_none=True), indent=4))


@evals.command()
@click.pass_context
@click.option(
    "--status",
    type=click.Choice(["pending", "queued", "running", "completed", "error", "user_error"]),
    help="Filter by job status.",
)
@click.option(
    "--limit",
    type=int,
    help="Limit number of results (max 100).",
)
def list(
    ctx: click.Context,
    status: Union[Literal["pending", "queued", "running", "completed", "error", "user_error"], None],
    limit: Union[int, None],
) -> None:
    """List evals"""

    client: Together = ctx.obj

    response = client.evals.list(status=status or omit, limit=limit or omit)

    display_list: List[Dict[str, Any]] = []
    for job in response:
        if job.parameters:
            model = job.parameters.get("model_to_evaluate", "")
            model_a = job.parameters.get("model_a", "")
            model_b = job.parameters.get("model_b", "")
        else:
            model = ""
            model_a = ""
            model_b = ""

        display_list.append(
            {
                "Workflow ID": job.workflow_id or "",
                "Type": job.type,
                "Status": job.status,
                "Created At": job.created_at or 0,
                "Model": model,
                "Model A": model_a,
                "Model B": model_b,
            }
        )

    table = tabulate(display_list, headers="keys", tablefmt="grid", showindex=True)
    click.echo(table)


@evals.command()
@click.pass_context
@click.argument("evaluation_id", type=str, required=True)
def retrieve(ctx: click.Context, evaluation_id: str) -> None:
    """Get details of a specific evaluation job"""

    client: Together = ctx.obj

    response = client.evals.retrieve(evaluation_id)

    click.echo(json.dumps(response.model_dump(exclude_none=True), default=datetime_serializer, indent=4))


@evals.command()
@click.pass_context
@click.argument("evaluation_id", type=str, required=True)
def status(ctx: click.Context, evaluation_id: str) -> None:
    """Get the status and results of a specific evaluation job"""

    client: Together = ctx.obj

    response = client.evals.status(evaluation_id)

    click.echo(json.dumps(response.model_dump(exclude_none=True), indent=4))


def _build_judge(
    type: Literal["classify", "score", "compare"],
    judge_model: str,
    judge_model_source: Literal["serverless", "dedicated", "external"],
    judge_system_template: str,
    judge_external_api_token: Optional[str],
    judge_external_base_url: Optional[str],
) -> ParametersEvaluationClassifyParametersJudge:
    if type == "classify":
        judge_config = ParametersEvaluationClassifyParametersJudge(
            model=judge_model,
            model_source=judge_model_source,
            system_template=judge_system_template,
        )
    elif type == "score":
        judge_config = ParametersEvaluationScoreParametersJudge(
            model=judge_model,
            model_source=judge_model_source,
            system_template=judge_system_template,
        )
    elif type == "compare":
        judge_config = ParametersEvaluationCompareParametersJudge(
            model=judge_model,
            model_source=judge_model_source,
            system_template=judge_system_template,
        )

    if judge_external_api_token:
        judge_config["external_api_token"] = judge_external_api_token
    if judge_external_base_url:
        judge_config["external_base_url"] = judge_external_base_url

    return judge_config
