import json
from typing import Any, Dict, Union, Literal, Optional, cast

import click

from together import Together, TogetherError
from together.lib.cli.api._utils import handle_api_errors
from together.types.eval_create_params import (
    ParametersEvaluationScoreParameters,
    ParametersEvaluationCompareParameters,
    ParametersEvaluationClassifyParameters,
    ParametersEvaluationScoreParametersJudge,
    ParametersEvaluationCompareParametersJudge,
    ParametersEvaluationClassifyParametersJudge,
    ParametersEvaluationScoreParametersModelToEvaluate,
    ParametersEvaluationClassifyParametersModelToEvaluate,
    ParametersEvaluationCompareParametersModelAEvaluationModelRequest,
    ParametersEvaluationCompareParametersModelBEvaluationModelRequest,
)


@click.command()
@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.",
)
@click.pass_context
@handle_api_errors("Evals")
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))


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
