import streamlit as st
from clarifai_grpc.channel.clarifai_channel import ClarifaiChannel
from clarifai_grpc.grpc.api import resources_pb2, service_pb2, service_pb2_grpc
from clarifai_grpc.grpc.api.status import status_code_pb2

# Function to make API call and get text completion
def get_text_completion(raw_text):
    PAT = '9b209aadda08410caf0b6d815b57e080'
    USER_ID = 'openai'
    APP_ID = 'chat-completion'
    MODEL_ID = 'GPT-4'
    MODEL_VERSION_ID = '5d7a50b44aec4a01a9c492c5a5fcf387'
    
    channel = ClarifaiChannel.get_grpc_channel()
    stub = service_pb2_grpc.V2Stub(channel)
    metadata = (('authorization', 'Key ' + PAT),)
    
    userDataObject = resources_pb2.UserAppIDSet(user_id=USER_ID, app_id=APP_ID)
    
    post_model_outputs_response = stub.PostModelOutputs(
        service_pb2.PostModelOutputsRequest(
            user_app_id=userDataObject,
            model_id=MODEL_ID,
            version_id=MODEL_VERSION_ID,
            inputs=[
                resources_pb2.Input(
                    data=resources_pb2.Data(
                        text=resources_pb2.Text(
                            raw=raw_text
                        )
                    )
                )
            ]
        ),
        metadata=metadata
    )
    
    if post_model_outputs_response.status.code != status_code_pb2.SUCCESS:
        raise Exception(f"Post model outputs failed, status: {post_model_outputs_response.status.description}")
    
    output = post_model_outputs_response.outputs[0]
    
    return output.data.text.raw

# Streamlit app
def main():
    st.title("Text Completion App")
    
    # Get user input
    raw_text = st.text_area("Enter text:", "i need to know about narcotics pemishment")
    
    # Perform text completion when button is clicked
    if st.button("Complete Text"):
        completion_result = get_text_completion(raw_text)
        st.success("Completion:\n{}".format(completion_result))

if __name__ == "__main__":
    main()