import os
import uuid
import streamlit as st
import boto3
from botocore.config import Config
from botocore.exceptions import ClientError

# ====== Environment-driven config (set these in HF Spaces → Settings → Secrets) ======
REGION = os.getenv("AWS_REGION", "us-east-1")
AGENT_ID = os.getenv("BEDROCK_AGENT_ID", "S3XZUUJNMN")           # Provided by you
AGENT_ALIAS_ID = os.getenv("BEDROCK_AGENT_ALIAS_ID", "JRECZBFSS5")# Provided by you

# ====== Streamlit page config ======
st.set_page_config(page_title="Bedrock Agent Chat", page_icon="🤖")
st.title("🤖 Bedrock Agent Chat")
st.caption("Type below and your Amazon Bedrock Agent will respond.")

# ====== Initialize session state ======
if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())

if "history" not in st.session_state:
    st.session_state.history = []

# ====== Helper: create Bedrock Agents Runtime client ======
def get_bedrock_agents_client():
    # Boto3 will automatically use the AWS credentials provided via environment variables
    # (on Hugging Face Spaces, set these in Settings → Secrets)
    return boto3.client(
        "bedrock-agent-runtime",
        region_name=REGION,
        config=Config(retries={"max_attempts": 3})
    )

# ====== Render chat history ======
for role, msg in st.session_state.history:
    with st.chat_message(role):
        st.markdown(msg)

# ====== Chat input box ======
user_text = st.chat_input("Type your message")
if user_text:
    st.session_state.history.append(("user", user_text))
    with st.chat_message("user"):
        st.markdown(user_text)

    with st.chat_message("assistant"):
        placeholder = st.empty()
        running_text = ""

        try:
            client = get_bedrock_agents_client()

            # Invoke the Agent (streaming)
            response = client.invoke_agent(
                agentId=AGENT_ID,
                agentAliasId=AGENT_ALIAS_ID,
                sessionId=st.session_state.session_id,
                inputText=user_text
                # Optionally add enableTrace=True for debugging/tracing
            )

            # Stream the completion events; final text arrives as chunk.bytes per API docs
            if "completion" in response and response["completion"] is not None:
                for event in response["completion"]:
                    if "chunk" in event and "bytes" in event["chunk"]:
                        piece = event["chunk"]["bytes"].decode("utf-8", errors="ignore")
                        running_text += piece
                        placeholder.markdown(running_text)

            if not running_text:
                running_text = "_(No response received from agent)_"
            st.session_state.history.append(("assistant", running_text))

        except ClientError as e:
            err = f"⚠️ AWS ClientError: {e.response.get('Error', {}).get('Message', str(e))}"
            st.error(err)
            st.session_state.history.append(("assistant", err))
        except Exception as e:
            err = f"⚠️ Unexpected error: {str(e)}"
            st.error(err)
            st.session_state.history.append(("assistant", err))