from typing import List
import json
import time
import os
from urllib import request
import sys


BASE_LEARNING_URL = "https://api.inworld.ai/learning/v1alpha/"
BASE_LRO_URL = "https://api.inworld.ai/lro/v1alpha/"


_EVAL_CONFIG = """
# List of evaluators to run
evaluators:

  # Pointwise evaluator: use LLM with custom prompt
  - name: engagement
    evaluator_type: pointwise
    args:
      model_name: gpt-4o-mini

      # Explicitly set the allowed labels to be used for metrics names
      labels:

        - CONTENT_ENGAGING
        - CONTENT_BORING

      # ">" is to eliminate newlines, {{input}} and {{output}} corresponds to input and output of evaluation node
      prompt_template: >
        A response must meet very specific criteria to be labeled CONTENT_ENGAGING.
        The bar for engagement should be high - most casual or mildly interesting responses
        should be labeled as CONTENT_BORING. Only exceptional content that creates
        meaningful entertainment value deserves the CONTENT_ENGAGING label.

        ## User Input
        {{ input }}

        ## Response to review
        {{ output }}

        Your output must ONLY be a JSON object. Nothing before or after. Your output should be structured as follow:
        {
            "result": "CONTENT_ENGAGING" | "CONTENT_BORING",
            "reason": "Detailed explanation of your assessment"
        }

  # Rule-based evaluator example: use jinja logical expressions
  - name: basic_rules
    evaluator_type: rules
    args:
      # This currently should be set to False to work
      json_parse: false
      logical_expressions:
        # Check output length
        - output|length > 5
        # Check whether all elements have required fields
        - output|selectattr('handle')|length == output|length and output|selectattr('content')|length == output|length
        # Check if any content contains a question mark (indicating engagement)
        - output|selectattr('content', 'contains', '?')|length > 0
"""


def run_evaluation(input_node_name: str, output_node_name: str, execution_ids: List[str], evaluation_group_name: str, token: str):
    result = _run_evaluation(input_node_name, output_node_name, execution_ids, evaluation_group_name, token)
    _poll_for_evaluation_run(result['name'], token)
    result = _list_evaluations(token)

    run_id = None
    for run_data in result['evaluations']:
        if run_data['evaluationId'] != evaluation_group_name:
            continue
        run_id = run_data['runName']
        _get_metrics(run_data['runName'], evaluation_group_name, token)
        break

    if run_id is None:
        print("Evaluation run is not complete or had some error.")
    else:
        print(f"Evaluation run is complete. Please check the Portal dashboard under evaluation group: {evaluation_group_name}, run_id: {run_id}")

def _run_evaluation(input_node_name: str, output_node_name: str, execution_ids: List[str], evaluation_group: str, token: str):
    payload = {
        "input_node_name": input_node_name,
        "output_node_name": output_node_name,
        "execution_ids": [{"id": x} for x in execution_ids],
        "evaluation_id": evaluation_group,
        "eval_config": _EVAL_CONFIG,
    }
    result = http_request(
        f"{BASE_LEARNING_URL}evaluations",
        token,
        method="POST",
        data=payload
    )
    return result

def _poll_for_evaluation_run(operations_path: str, token: str):
    done = False
    count = 0
    while not done:
        response_json = http_request(
            f"{BASE_LRO_URL}{operations_path}",
            token,
        )
        done = response_json.get('done', False)
        if not done:
            count += 1
            time.sleep(1)

        if count > 180:
            break

def _list_evaluations(token: str):
    return http_request(
        f"{BASE_LEARNING_URL}evaluations",
        token,
        method="GET"
    )

def _get_metrics(run_id: str, evaluation_group: str, token: str):
    result = http_request(
        f"{BASE_LEARNING_URL}evaluations/{evaluation_group}/runs/{run_id}/metrics",
        token,
        method="GET"
    )
    return result

def http_request(url, token, method="GET", headers=None, data=None):
    if headers is None:
        headers = {}
    headers["Authorization"] = f"Basic {token}"
    if method == "POST":
        if data is not None:
            data = json.dumps(data).encode()
        req = request.Request(url, data=data, headers={**headers, "Content-Type": "application/json"}, method=method)
    else:
        req = request.Request(url, headers=headers, method=method)

    try:
        with request.urlopen(req) as response:
            raw_resp = response.read().decode()
            try:
                return json.loads(raw_resp)
            except Exception:
                return raw_resp
    except request.HTTPError as response:
        if response.status == 409:
            return json.loads(response.read().decode())
        else:
            raise Exception(f"Failed to find traces for evaluation run: {response}")

if __name__ == "__main__":
    inps = sys.stdin.readline().strip()
    args = inps.split(" ")
    if len(args) != 5:
        print(
            "Invalid number of arguments. Need input and output node names, execution ids, "
              "an evaluation group name, and a base64 auth token"
              )
        sys.exit(1)
    try:
        run_evaluation(args[0], args[1], args[2].split(","), args[3], args[4])
    except Exception as e:
        print(str(e))
