evaluator.py

import asyncio
import json
import logging
import math
from typing import Dict, Any, List, Optional, Tuple, Union

import aiohttp
from jsonschema.exceptions import ValidationError

from config import logger, DEFAULT_REQUEST_DELAY, DEFAULT_MAX_CONCURRENCY, DEFAULT_REQUEST_TIMEOUT, MAX_ERROR_THRESHOLD, MAX_RETRY_ATTEMPTS
from mued_schema import MuEdSchema, SchemaLoadError, REQUIRED_ARTEFACT_TYPE, get_schema as get_mued_schema

LAMBDA_INVOKE_PATH = "/2015-03-31/functions/function/invocations"
TRANSPORTS = ("http", "lambda")


def _prepare_legacy_payload(record: Dict[str, Any], eval_function_params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """Constructs the legacy Lambda-Feedback JSON payload for the API request from the DB record."""
    grade_params = record.get('grade_params') or {}
    response = record.get('submission')
    answer = record.get('answer').replace('"', '')

    logging.debug(f"Response Type: {response} -  {type(response)}")
    logging.debug(f"Answer Type: {answer} -  {type(answer)}")

    cases = [
        {**case, "params": case["params"] if case.get("params") is not None else {}}
        for case in record.get('cases', []) or []
    ]

    symbols = record.get('symbols') or {}

    return {
        "response": response,
        "answer": answer,
        "params": {
            **grade_params,
            **(eval_function_params or {}),
            "symbols": symbols,
            "cases": cases,
        }
    }


def _prepare_mued_payload(record: Dict[str, Any], schema: MuEdSchema) -> Dict[str, Any]:
    """Constructs a muEd EvaluateRequest JSON payload from the DB record, validated against the
    live muEd OpenAPI schema before being returned."""
    response = record.get('submission')

    payload: Dict[str, Any] = {
        "submission": {
            "type": REQUIRED_ARTEFACT_TYPE,
            "content": {"value": response},
        },
    }

    answer = record.get('answer')
    if answer is not None:
        if isinstance(answer, str):
            answer = answer.replace('"', '')
        payload["task"] = {
            "title": "Evaluation Task",
            "referenceSolution": {"value": answer},
        }

    grade_params = record.get('grade_params') or {}
    cases = [
        {**case, "params": case["params"] if case.get("params") is not None else {}}
        for case in record.get('cases', []) or []
    ]
    symbols = record.get('symbols') or {}
    config_params = {**grade_params, "symbols": symbols, "cases": cases}
    if config_params:
        payload["configuration"] = {"params": config_params}

    schema.validate_evaluate_request(payload)
    return payload


def _wrap_lambda_event(payload: Dict[str, Any], path: str) -> Dict[str, Any]:
    """Wraps a request payload in an API Gateway proxy event, as a Lambda invocation (e.g. via the
    AWS Runtime Interface Emulator) expects. Both v1 (path/httpMethod) and v2 (rawPath/requestContext.http)
    routing fields are set so either handler style routes correctly."""
    return {
        "body": json.dumps(payload),
        "path": path,
        "rawPath": path,
        "httpMethod": "POST",
        "requestContext": {"http": {"method": "POST", "path": path}},
        "headers": {"content-type": "application/json"},
        "isBase64Encoded": False,
    }


def _http_error(status_code: Any, text: str) -> Dict[str, Any]:
    """Builds an "HTTP Error" network error. A JSON object body is kept whole under response_json so
    callers can inspect structured error responses (response_text is truncated)."""
    error: Dict[str, Any] = {
        "error_type": "HTTP Error",
        "status_code": status_code,
        "message": f"Received status code {status_code}.",
        "response_text": text[:200],
    }
    try:
        parsed = json.loads(text)
    except ValueError:
        parsed = None
    if isinstance(parsed, dict):
        error["response_json"] = parsed
    return error


def _mued_error_to_grader_exception(execution_error: Dict[str, Any], db_grade: Any) -> Optional[Dict[str, Any]]:
    """A muEd service reports grading failures as HTTP 500 with an ErrorResponse body ({title, message?, code?}),
    where the legacy API returned 200 with an 'error' object. Converts such an error into a "Grader Exception"
    so it's counted and analysed like a legacy one. Other statuses (400 = our payload rejected, 503, proxy errors
    without an ErrorResponse body) stay network errors. Returns None if not applicable."""
    if execution_error.get("error_type") != "HTTP Error" or execution_error.get("status_code") != 500:
        return None
    body = execution_error.get("response_json")
    if not isinstance(body, dict) or not isinstance(body.get("title"), str):
        return None

    message = body.get("message") or body["title"]
    return {
        "error_type": "Grader Exception",
        "message": message,
        "detail": message,
        "code": body.get("code"),
        "status_code": 500,
        "original_grade": db_grade,
    }


def _unwrap_lambda_response(data: Any) -> Tuple[Optional[Any], Optional[Dict[str, Any]]]:
    """Unwraps a Lambda invocation result. API Gateway proxy envelopes ({statusCode, body}) are decoded
    like an HTTP response; unhandled handler crashes ({errorMessage, errorType}) become an error; anything
    else is returned unchanged. Returns (response_data, error_details)."""
    if not isinstance(data, dict):
        return data, None

    if "statusCode" in data:
        body = data.get("body") or ""
        if data["statusCode"] != 200:
            return None, _http_error(data["statusCode"], body if isinstance(body, str) else json.dumps(body))
        if not isinstance(body, str):
            return body, None
        try:
            return json.loads(body), None
        except ValueError:
            return None, {
                "error_type": "JSON Decode Error",
                "message": "Lambda response body could not be parsed as JSON.",
                "response_text": body[:200]
            }

    if "errorMessage" in data:
        return None, {
            "error_type": "Lambda Invocation Error",
            "message": str(data.get("errorMessage"))[:500],
            "lambda_error_type": data.get("errorType"),
        }

    return data, None


async def _execute_request(session: aiohttp.ClientSession, endpoint_path: str, payload: Dict[str, Any], retry_base_delay: float = 0.0,
                           lambda_path: Optional[str] = None,
                           request_timeout: float = DEFAULT_REQUEST_TIMEOUT) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
    """Executes the POST request with exponential backoff on transient network errors.
    When lambda_path is set, the payload is sent as an API Gateway proxy event for that path and the
    Lambda result is unwrapped. Timeouts are not retried in that case: the invocation may still be running,
    and a second invocation against a single-slot runtime (e.g. the AWS RIE) crashes it.
    Returns (response_data, error_details).
    """
    last_error: Optional[Exception] = None
    body = _wrap_lambda_event(payload, lambda_path) if lambda_path else payload

    for attempt in range(MAX_RETRY_ATTEMPTS + 1):
        try:
            async with session.post(
                endpoint_path,
                json=body,
                timeout=aiohttp.ClientTimeout(total=request_timeout),
            ) as response:
                if response.status != 200:
                    return None, _http_error(response.status, await response.text())

                try:
                    data = await response.json(content_type=None)
                except Exception:
                    text = await response.text()
                    # The AWS RIE reports a function timeout as a plain-text 200, not JSON.
                    if lambda_path and text.startswith("Task timed out"):
                        return None, {"error_type": "Timeout", "message": text[:200]}
                    return None, {
                        "error_type": "JSON Decode Error",
                        "message": "API response could not be parsed as JSON.",
                        "response_text": text[:200]
                    }

                return _unwrap_lambda_response(data) if lambda_path else (data, None)

        except (aiohttp.ClientError, asyncio.TimeoutError) as e:
            last_error = e
            if lambda_path and isinstance(e, asyncio.TimeoutError):
                return None, {
                    "error_type": "Timeout",
                    "message": f"No response within {request_timeout}s; not retried (Lambda invocation may still be running).",
                }
            if attempt < MAX_RETRY_ATTEMPTS:
                wait = max(retry_base_delay, 1.0) * (2 ** attempt)
                logger.warning(
                    f"Network error on attempt {attempt + 1}/{MAX_RETRY_ATTEMPTS + 1} "
                    f"({type(e).__name__}: {e}). Retrying in {wait:.1f}s..."
                )
                await asyncio.sleep(wait)

    return None, {
        "error_type": "ConnectionError",
        "message": f"{type(last_error).__name__}: {last_error}",
        "retries_attempted": MAX_RETRY_ATTEMPTS,
    }


def _validate_response(response_data: Dict[str, Any], db_grade: Any) -> Optional[Dict[str, Any]]:
    """Compares the API's 'is_correct' result against the historical database grade."""
    result = response_data.get('result')

    if result is None:
        api_error = response_data.get('error', {})
        return {
            "error_type": "Grader Exception",
            "message": api_error.get('message', 'API returned no result.'),
            "detail": api_error.get('detail', ''),
            "original_grade": db_grade,
        }

    api_is_correct = result.get('is_correct')

    expected_is_correct: Optional[bool]
    if isinstance(db_grade, int):
        expected_is_correct = bool(db_grade)
    elif db_grade is None:
        expected_is_correct = None
    else:
        expected_is_correct = db_grade

    if api_is_correct is None:
        return {
            "error_type": "Missing API Field",
            "message": "API response is missing the 'is_correct' field.",
            "original_grade": db_grade
        }

    if api_is_correct == expected_is_correct:
        return None

    return {
        "error_type": "**Grade Mismatch**",
        "message": f"API result '{api_is_correct}' does not match DB grade '{expected_is_correct}'.",
        "original_grade": db_grade
    }


def _validate_mued_response(response_data: Any, db_grade: Any) -> Optional[Dict[str, Any]]:
    """Compares a muEd Feedback[] response's awardedPoints against the historical database grade.
    awardedPoints == 1 is treated as correct, == 0 as incorrect; anything else is a reported mismatch."""
    if not isinstance(response_data, list):
        return {
            "error_type": "Malformed muEd Response",
            "message": f"Expected a JSON array of Feedback items, got {type(response_data).__name__}.",
            "original_grade": db_grade,
        }

    if not response_data:
        return {
            "error_type": "Empty muEd Response",
            "message": "API returned an empty Feedback[] array; cannot determine awardedPoints.",
            "original_grade": db_grade,
        }

    awarded_points = response_data[0].get('awardedPoints')

    expected_is_correct: Optional[bool]
    if isinstance(db_grade, int):
        expected_is_correct = bool(db_grade)
    elif db_grade is None:
        expected_is_correct = None
    else:
        expected_is_correct = db_grade

    if awarded_points not in (0, 1):
        return {
            "error_type": "Unexpected awardedPoints",
            "message": f"awardedPoints={awarded_points!r}; expected 0 or 1 for pass/fail mapping.",
            "original_grade": db_grade,
        }

    api_is_correct = bool(awarded_points)

    if api_is_correct == expected_is_correct:
        return None

    return {
        "error_type": "**Grade Mismatch**",
        "message": f"muEd awardedPoints={awarded_points} (-> {api_is_correct}) does not match DB grade '{expected_is_correct}'.",
        "original_grade": db_grade,
    }


def _check_feedback(response_data: Dict[str, Any], db_feedback: Any) -> Optional[Dict[str, Any]]:
    """Checks if API feedback matches the stored DB feedback. Returns a warning dict or None."""
    result = response_data.get('result', {})
    api_feedback = result.get('feedback')

    if db_feedback is None or api_feedback is None:
        return None

    if api_feedback != db_feedback:
        return {
            "warning_type": "Feedback Mismatch",
            "message": "API feedback does not match DB feedback.",
            "db_feedback": db_feedback,
            "api_feedback": api_feedback,
        }
    return None


async def test_endpoint(base_endpoint: str, data_records: List[Dict[str, Any]],
                        request_delay: float = DEFAULT_REQUEST_DELAY,
                        max_concurrency: int = DEFAULT_MAX_CONCURRENCY,
                        mode: str = "legacy",
                        mued_schema_path: Optional[str] = None,
                        max_error_threshold: Union[int, float] = MAX_ERROR_THRESHOLD,
                        eval_function_params: Optional[Dict[str, Any]] = None,
                        transport: str = "http",
                        request_timeout: float = DEFAULT_REQUEST_TIMEOUT) -> Dict[str, Any]:
    """Tests the endpoint against all records concurrently, returns aggregated results.

    transport "http" POSTs payloads directly to the endpoint; "lambda" treats the endpoint as a Lambda
    invoke URL (e.g. the AWS Runtime Interface Emulator) and sends each payload as an API Gateway proxy event.

    max_error_threshold is either a whole number (absolute error count) or a float in
    [0.0, 1.0] (fraction of total_records) at which either the validation-error or
    network-error count triggers an early stop.
    """
    total_records = len(data_records)
    successful_requests = 0
    errors = []
    network_errors = []
    feedback_warnings = []
    parsing_warnings = []
    validation_error_count = 0
    network_error_count = 0
    completed_count = 0

    if transport not in TRANSPORTS:
        raise ValueError(f"transport must be one of {TRANSPORTS} (got {transport!r})")

    if mode == "mued":
        schema = get_mued_schema(mued_schema_path) if mued_schema_path else get_mued_schema()
        request_path = base_endpoint.rstrip('/') + '/evaluate'
        route_path = '/evaluate'
        prepare_fn = lambda record: _prepare_mued_payload(record, schema)
        validate_fn = _validate_mued_response
        check_feedback_fn = None
    else:
        request_path = base_endpoint
        route_path = '/'
        prepare_fn = lambda record: _prepare_legacy_payload(record, eval_function_params)
        validate_fn = _validate_response
        check_feedback_fn = _check_feedback

    lambda_path = None
    if transport == "lambda":
        lambda_path = route_path
        request_path = base_endpoint if base_endpoint.rstrip('/').endswith('/invocations') \
            else base_endpoint.rstrip('/') + LAMBDA_INVOKE_PATH

    if isinstance(max_error_threshold, float):
        if not 0.0 <= max_error_threshold <= 1.0:
            raise ValueError(f"max_error_threshold as a float must be between 0.0 and 1.0 (got {max_error_threshold})")
        effective_error_threshold = max(1, math.ceil(max_error_threshold * total_records))
        logger.info(f"Error threshold: {max_error_threshold:.0%} of {total_records} records ({effective_error_threshold} errors)")
    else:
        effective_error_threshold = max_error_threshold

    semaphore = asyncio.Semaphore(max_concurrency)
    lock = asyncio.Lock()
    stop_event = asyncio.Event()

    logger.info(f"Starting tests on endpoint with mode={mode}, transport={transport}, request_timeout={request_timeout}s, max_concurrency={max_concurrency}, request_delay={request_delay}s")

    async def worker(i: int, record: Dict[str, Any]) -> None:
        nonlocal validation_error_count, network_error_count, completed_count, successful_requests

        if stop_event.is_set():
            return

        async with semaphore:
            if stop_event.is_set():
                return

            submission_id = str(record.get('submission_id')) if record.get('submission_id') is not None else None

            try:
                payload = prepare_fn(record)
            except (SchemaLoadError, ValidationError) as e:
                async with lock:
                    validation_error_count += 1
                    errors.append({
                        "error_type": "Schema Validation Error",
                        "message": str(e),
                        "original_grade": record.get('grade'),
                        "submission_id": submission_id,
                    })
                    if validation_error_count >= MAX_ERROR_THRESHOLD:
                        logger.warning(f"Stopping early! Reached maximum error threshold of {MAX_ERROR_THRESHOLD}.")
                        stop_event.set()
                    completed_count += 1
                return

            response_data, execution_error = await _execute_request(session, request_path, payload, retry_base_delay=request_delay,
                                                                    lambda_path=lambda_path, request_timeout=request_timeout)
            logging.debug(f"[{submission_id}] grade={record.get('grade')} | REQUEST: {payload} | RESPONSE: {response_data or execution_error}")

            grader_exception = None
            if execution_error and mode == "mued":
                grader_exception = _mued_error_to_grader_exception(execution_error, record.get('grade'))
                if grader_exception:
                    execution_error = None

            async with lock:
                if execution_error:
                    network_error_count += 1
                    execution_error['submission_id'] = submission_id
                    execution_error['original_grade'] = record.get('grade')
                    execution_error['request_payload'] = payload
                    network_errors.append(execution_error)
                    if network_error_count >= effective_error_threshold:
                        logger.warning(f"Stopping early! Reached maximum error threshold of {effective_error_threshold}.")
                        stop_event.set()
                    completed_count += 1
                    if completed_count % 10 == 0:
                        logger.info(f"Progress: {completed_count}/{total_records} requests completed")
                    return

                validation_error = grader_exception or validate_fn(response_data, record.get('grade'))

                if validation_error:
                    if (record.get('historical_error_message') is not None
                            and validation_error.get('error_type') == 'Grader Exception'):
                        parsing_warnings.append({
                            "warning_type": "Parsing Error",
                            "message": validation_error.get('message'),
                            "detail": validation_error.get('detail', ''),
                            "historical_error_message": record.get('historical_error_message'),
                            "historical_error_detail": record.get('historical_error_detail'),
                            "submission_id": submission_id,
                            "request_payload": payload,
                        })
                        successful_requests += 1
                    else:
                        validation_error_count += 1
                        validation_error['submission_id'] = submission_id
                        validation_error['request_payload'] = payload
                        errors.append(validation_error)
                        if validation_error_count >= effective_error_threshold:
                            logger.warning(f"Stopping early! Reached maximum error threshold of {effective_error_threshold}.")
                            stop_event.set()
                else:
                    successful_requests += 1

                if check_feedback_fn:
                    feedback_warning = check_feedback_fn(response_data, record.get('feedback'))
                    if feedback_warning:
                        feedback_warning['submission_id'] = submission_id
                        feedback_warning['request_payload'] = payload
                        feedback_warnings.append(feedback_warning)

                completed_count += 1
                if completed_count % 10 == 0:
                    logger.info(f"Progress: {completed_count}/{total_records} requests completed")

    connector = aiohttp.TCPConnector(limit=max_concurrency + 2)
    async with aiohttp.ClientSession(connector=connector) as session:
        tasks = []
        for i, record in enumerate(data_records):
            if stop_event.is_set():
                break
            tasks.append(asyncio.create_task(worker(i, record)))
            if request_delay > 0 and i < total_records - 1:
                await asyncio.sleep(request_delay)
        await asyncio.gather(*tasks, return_exceptions=True)

    return {
        "pass_count": successful_requests,
        "total_count": total_records,
        "tested_count": completed_count,
        "number_of_errors": validation_error_count,
        "number_of_network_errors": network_error_count,
        "list_of_errors": errors,
        "list_of_network_errors": network_errors,
        "list_of_feedback_warnings": feedback_warnings,
        "list_of_parsing_warnings": parsing_warnings,
    }