import json
import boto3
import uuid

client = boto3.client("bedrock-agentcore")

RUNTIME_ARN = (
    "arn:aws:bedrock-agentcore:us-east-1:017485772833:"
    "runtime/nmpOnboardingAgent_NmpOnboardingAgent-Azp1olHW9k"
)


def extract_text_from_sse(sse_text: str) -> str:
    parts = []

    for line in sse_text.splitlines():
        if not line.startswith("data: "):
            continue

        try:
            event = json.loads(line[6:])
        except json.JSONDecodeError:
            continue

        delta = (
            event.get("event", {})
            .get("contentBlockDelta", {})
            .get("delta", {})
        )

        text = delta.get("text")

        if text:
            parts.append(text)

    return "".join(parts)


def response(status_code, body):
    return {
        "statusCode": status_code,
        "headers": {
            "Content-Type": "application/json",
            "Access-Control-Allow-Origin": "*"
        },
        "body": json.dumps(body, ensure_ascii=False)
    }


def lambda_handler(event, context):
    try:
        body = event.get("body", {})

        if isinstance(body, str):
            body = json.loads(body)

        prompt = body.get("prompt", "").strip()

        if not prompt:
            return response(
                400,
                {
                    "error": "El campo prompt es obligatorio"
                }
            )

        # Si el kiosco manda sessionId, continuamos la conversación.
        # Si no lo manda, iniciamos una nueva.
        session_id = body.get("sessionId")

        if not session_id:
            session_id = str(uuid.uuid4())

        payload = {
            "prompt": prompt
        }

        agent_response = client.invoke_agent_runtime(
            agentRuntimeArn=RUNTIME_ARN,
            runtimeSessionId=session_id,
            payload=json.dumps(payload).encode("utf-8")
        )

        raw_response = (
            agent_response["response"]
            .read()
            .decode("utf-8")
        )

        final_text = extract_text_from_sse(raw_response)

        return response(
            200,
            {
                "sessionId": session_id,
                "response": final_text
            }
        )

    except json.JSONDecodeError:
        return response(
            400,
            {
                "error": "El body no contiene JSON valido"
            }
        )

    except Exception as e:
        print(f"ERROR: {str(e)}")

        return response(
            500,
            {
                "error": "Error al invocar el agente"
            }
        )