import json
import logging
import os

import boto3
from botocore.config import Config
from botocore.exceptions import BotoCoreError, ClientError, ReadTimeoutError

# Keep deployment-specific settings configurable.
MODEL_ID = os.environ.get(
"MODEL_ID",
"global.anthropic.claude-haiku-5-5",
)

BEDROCK_REGION = os.environ.get(
"BEDROCK_REGION",
os.environ.get("AWS_REGION", "us-east-1"),
)

MAX_INPUT_CHARS = int(
os.environ.get("MAX_INPUT_CHARS", "8000")
)

MAX_OUTPUT_TOKENS = int(
os.environ.get("MAX_OUTPUT_TOKENS", "512")
)

logger = logging.getLogger()
logger.setLevel(logging.INFO)

# Best practice: initialize the SDK client outside the handler.
# Warm Lambda invocations can reuse the client and its connections.
bedrock = boto3.client(
"bedrock-runtime",
region_name=BEDROCK_REGION,
config=Config(
connect_timeout=3,
read_timeout=20,
# One SDK attempt avoids retrying a slow request beyond
# the intended Lambda execution window.
retries={
"total_max_attempts": 1,
"mode": "standard",
},
),
)

# Security measure: expose a controlled set of operations.
# Clients select an operation but cannot supply arbitrary system prompts.
SYSTEM_PROMPTS = {
"summarize": (
"You are a precise summarizer. "
"Summarize the supplied text in at most three sentences."
),
"sentiment": (
"Classify the sentiment of the supplied text as positive, "
"negative, or neutral. Reply with one word only."
),
"keywords": (
"Extract up to eight relevant keywords from the supplied text. "
"Return a comma-separated list only."
),
}

def make_response(status_code, body):
"""Create a standard API Gateway proxy response."""
return {
"statusCode": status_code,
"headers": {
"Content-Type": "application/json",
"Cache-Control": "no-store",
},
"body": json.dumps(body),
}

def lambda_handler(event, context):
# Step 1: parse and validate the request before using the model.
try:
raw_body = event.get("body") or "{}"

if isinstance(raw_body, str):
payload = json.loads(raw_body)
elif isinstance(raw_body, dict):
# Also supports direct invocation and unit tests.
payload = raw_body
else:
return make_response(
400,
{"error": "Request body must be a JSON object."},
)

except json.JSONDecodeError:
return make_response(
400,
{"error": "Body must contain valid JSON."},
)

if not isinstance(payload, dict):
return make_response(
400,
{"error": "Request body must be a JSON object."},
)

task = payload.get("task", "summarize")
text = payload.get("text", "")

# Reject unknown operations before an inference request is made.
if not isinstance(task, str) or task not in SYSTEM_PROMPTS:
return make_response(
400,
{
"error": (
f"Task must be one of {sorted(SYSTEM_PROMPTS)}"
)
},
)

if not isinstance(text, str) or not text.strip():
return make_response(
400,
{"error": "A non-empty text field is required."},
)

text = text.strip()

# Bound request size to reduce excessive input consumption.
# Character count is an application guardrail, not a token count.
if len(text) > MAX_INPUT_CHARS:
return make_response(
413,
{
"error": (
f"Text exceeds {MAX_INPUT_CHARS} characters."
)
},
)

# Step 2: invoke Claude Haiku 5.5 with the ConverseStream API.
try:
result = bedrock.converse_stream(
modelId=MODEL_ID,
system=[
{"text": SYSTEM_PROMPTS[task]}
],
messages=[
{
"role": "user",
"content": [
{"text": text}
],
}
],
inferenceConfig={
# Bound generated output to control cost and latency.
"maxTokens": MAX_OUTPUT_TOKENS,
"temperature": 0.2,
},
)

chunks = []
usage = {}

# Read Bedrock's event stream and collect text fragments.
# The client will receive a buffered JSON response below.
for stream_event in result["stream"]:
if "contentBlockDelta" in stream_event:
delta = stream_event["contentBlockDelta"].get(
"delta", {}
)
chunks.append(delta.get("text", ""))

elif "metadata" in stream_event:
usage = stream_event["metadata"].get("usage", {})

except ReadTimeoutError:
logger.warning(
"Bedrock read timeout: request_id=%s",
context.aws_request_id,
)

return make_response(
504,
{"error": "The model request timed out."},
)

except ClientError as error:
error_code = error.response.get(
"Error", {}
).get("Code", "Unknown")

# Log the service error code, not the user's input text.
logger.error(
"Bedrock invocation failed: code=%s request_id=%s",
error_code,
context.aws_request_id,
)

if error_code in {
"ThrottlingException",
"ServiceQuotaExceededException",
}:
return make_response(
429,
{"error": "The model is busy. Retry later."},
)

return make_response(
502,
{"error": "The upstream model request failed."},
)

except BotoCoreError:
logger.exception(
"Bedrock SDK failure: request_id=%s",
context.aws_request_id,
)

return make_response(
502,
{"error": "Unable to complete the model request."},
)

# Step 3: record usage so request costs can be analyzed.
# Never log the full prompt or generated content by default.
logger.info(
"request_id=%s task=%s input_tokens=%s output_tokens=%s",
context.aws_request_id,
task,
usage.get("inputTokens"),
usage.get("outputTokens"),
)

return make_response(
200,
{
"task": task,
"result": "".join(chunks).strip(),
"usage": usage,
"request_id": context.aws_request_id,
},
)