"""Job store client — DynamoDB.

Stores eval job status/results so `GET /run-eval/{job_id}` can be served by any replica,
independent of which one actually ran the job. The full result is kept as a single JSON string
attribute (`result_json`) rather than native DynamoDB Map/List types, since the result is only
ever read back whole (never queried into) and this avoids converting every float in it to
`Decimal`, which DynamoDB's native number type requires.
"""

import os
import time
from datetime import datetime, timezone
from functools import lru_cache
from typing import Any, NamedTuple

import boto3
from botocore.exceptions import ClientError

from ai_eval_runner import config


class JobRecord(NamedTuple):
    status: str
    result_json: str | None
    error: str | None


@lru_cache(maxsize=1)
def _dynamodb_resource() -> Any:
    # dynamodb-local doesn't validate the region, but boto3 still needs one to sign requests.
    region = os.environ.get("AWS_REGION", config.BEDROCK_REGION)
    return boto3.resource(
        "dynamodb", region_name=region, endpoint_url=config.DYNAMODB_ENDPOINT_URL
    )


def _table() -> Any:
    return _dynamodb_resource().Table(config.JOBS_TABLE_NAME)


def ensure_table_exists() -> None:
    """Create the jobs table against a local DynamoDB. Dev-only — never touches real AWS."""
    if not config.DYNAMODB_ENDPOINT_URL:
        return

    resource = _dynamodb_resource()
    try:
        resource.create_table(
            TableName=config.JOBS_TABLE_NAME,
            AttributeDefinitions=[{"AttributeName": "job_id", "AttributeType": "S"}],
            KeySchema=[{"AttributeName": "job_id", "KeyType": "HASH"}],
            BillingMode="PAY_PER_REQUEST",
        ).wait_until_exists()
    except ClientError as exc:
        if exc.response["Error"]["Code"] != "ResourceInUseException":
            raise

    resource.meta.client.update_time_to_live(
        TableName=config.JOBS_TABLE_NAME,
        TimeToLiveSpecification={"Enabled": True, "AttributeName": "expires_at"},
    )


def create_job(job_id: str, mcp_name: str, pipeline_id: str) -> None:
    now = datetime.now(timezone.utc).isoformat()
    _table().put_item(
        Item={
            "job_id": job_id,
            "status": "pending",
            "mcp_name": mcp_name,
            "pipeline_id": pipeline_id,
            "created_at": now,
            "updated_at": now,
            "expires_at": int(time.time()) + config.JOB_TTL_SECONDS,
        }
    )


def update_job(
    job_id: str,
    status: str,
    result_json: str | None = None,
    error: str | None = None,
) -> None:
    expression_names = {"#status": "status"}
    expression_values: dict[str, Any] = {
        ":status": status,
        ":updated_at": datetime.now(timezone.utc).isoformat(),
    }
    set_clauses = ["#status = :status", "updated_at = :updated_at"]

    if result_json is not None:
        set_clauses.append("result_json = :result_json")
        expression_values[":result_json"] = result_json
    if error is not None:
        set_clauses.append("#error = :error")
        expression_names["#error"] = "error"
        expression_values[":error"] = error

    _table().update_item(
        Key={"job_id": job_id},
        UpdateExpression="SET " + ", ".join(set_clauses),
        ExpressionAttributeNames=expression_names,
        ExpressionAttributeValues=expression_values,
    )


def get_job(job_id: str) -> JobRecord | None:
    item = _table().get_item(Key={"job_id": job_id}).get("Item")
    if item is None:
        return None
    return JobRecord(
        status=item["status"],
        result_json=item.get("result_json"),
        error=item.get("error"),
    )