import boto3
import json
import re


# -----------------------------
# JSON extraction
# -----------------------------

def extract_json_block(text: str, pick="last"):

    if not isinstance(text, str):
        return None

    cleaned = (
        text.replace("```json", "")
            .replace("```", "")
    )

    blocks = re.findall(r"\{[\s\S]*?\}", cleaned)

    if not blocks:
        return None

    block = blocks[-1] if pick == "last" else blocks[0]

    block = re.sub(r",\s*}", "}", block)

    try:
        return json.loads(block)
    except:
        return None


# -----------------------------
# Bedrock client
# -----------------------------

client = boto3.client(
    "bedrock-runtime",
    region_name="eu-central-1",
)

MODEL_ID = "eu.amazon.nova-pro-v1:0"


# -----------------------------
# 🔥 ONLY RESPONSIBILITY: CALL LLM
# -----------------------------

def call_llm(prompt: str):

    message_list = [
        {
            "role": "user",
            "content": [{"text": prompt}],
        }
    ]

    system_list = [
        {
            "text": "You are a precise medical extraction assistant."
        }
    ]

    native_request = {
        "schemaVersion": "messages-v1",
        "messages": message_list,
        "system": system_list,
        "inferenceConfig": {
            "maxTokens": 800,
            "topP": 0.1,
            "topK": 20,
            "temperature": 0.2,
        },
    }

    response = client.invoke_model(
        modelId=MODEL_ID,
        body=json.dumps(native_request)
    )

    model_response = json.loads(response["body"].read())

    return model_response["output"]["message"]["content"][0]["text"]