import gradio as gr
import random
import requests
import os
from datetime import datetime
import hashlib
import hmac
import json

def sign(key, msg):
    return hmac.new(key, msg.encode('utf-8'), hashlib.sha256).digest()

def get_signature_key(secret_key, date_stamp, region_name, service_name):
    k_date = sign(('AWS4' + secret_key).encode('utf-8'), date_stamp)
    k_region = sign(k_date, region_name)
    k_service = sign(k_region, service_name)
    k_signing = sign(k_service, 'aws4_request')
    return k_signing

def generate_response(input, history):
    #bot_message = random.choice(["How are you?", "I love you", "I'm very hungry"])

    #Authentication AWS
    access_key = os.environ['AccessKey']
    secret_key = os.environ['SecretKey']
    region = 'eu-west-1' # e.g., 'us-west-2'
    service = 'sagemaker'
    
    #Hit the Sagemaker API and get results
    
    url = "https://runtime.sagemaker.eu-west-1.amazonaws.com/endpoints/jumpstart-dft-meta-textgeneration-llama-2-7b-f/invocations"
    #authorization_token = os.environ['Authorization']
    # headers = {
    #     "Content-Type": "application/json",
    #     "X-Amz-Date": datetime.utcnow().strftime("%Y%m%dT%H%M%SZ"),  # Using ISO-8601 basic format
    #     "X-Amzn-SageMaker-Custom-Attributes": "accept_eula=true",
    #     "Connection": "keep-alive"
    # }
    payload = {
        "inputs": [
            [
                {
                    "role": "system",
                    "content": "Always answer positively"
                },
                {
                    "role": "user",
                    "content": input #"I am going to Paris, what should I see?"
                }
            ]
        ],
        "parameters": {
            "max_new_tokens": 512,
            "top_p": 0.9,
            "temperature": 0.6
        }
    }

    payload_str = json.dumps(payload)

    t = datetime.utcnow()
    amz_date = t.strftime('%Y%m%dT%H%M%SZ')
    date_stamp = t.strftime('%Y%m%d')

    endpoint = "runtime.sagemaker.eu-west-1.amazonaws.com"
    canonical_uri = '/endpoints/jumpstart-dft-meta-textgeneration-llama-2-7b-f/invocations' # Your canonical URI here
    canonical_headers = f'host:{endpoint}\nx-amz-date:{amz_date}\n'
    signed_headers = 'host;x-amz-date'
    
    hashlib_payload = hashlib.sha256(payload_str.encode('utf-8')).hexdigest()
    canonical_request = f'POST\n{canonical_uri}\n\n{canonical_headers}\n{signed_headers}\n{hashlib_payload}'
    
    algorithm = 'AWS4-HMAC-SHA256'
    credential_scope = f'{date_stamp}/{region}/{service}/aws4_request'
    string_to_sign = f'{algorithm}\n{amz_date}\n{credential_scope}\n{hashlib.sha256(canonical_request.encode("utf-8")).hexdigest()}'
    
    signing_key = get_signature_key(secret_key, date_stamp, region, service)
    signature = hmac.new(signing_key, (string_to_sign).encode('utf-8'), hashlib.sha256).hexdigest()
    
    authorization_header = f'{algorithm} Credential={access_key}/{credential_scope}, SignedHeaders={signed_headers}, Signature={signature}'
    
    headers = {
        'x-amz-date': amz_date,
        'Authorization': authorization_header,
        'Content-Type': 'application/json',
        "X-Amzn-SageMaker-Custom-Attributes": "accept_eula=true"
    }
    
    #response = requests.post(endpoint, headers=headers, data=payload_str)

    
    response = requests.post(url, json=payload, headers=headers)
    
    if response.status_code == 200:
        response_json = response.json()
        if response_json and 'generation' in response_json[0]:
            content = response_json[0]['generation']['content']
            print("Content:", content)
            bot_message = content
        else:
            print("Content not found in response")
    else:
        print("Request failed:", response.status_code, response.text)


    return bot_message
    
def my_chatbot(input, history):

    history = history or []

    my_history = list(sum(history, ()))

    my_history.append(input)

    my_input = ' '.join(my_history)

    output = generate_response(input, history)

    history.append((input, output))

    return history, history

with gr.Blocks() as demo:

    gr.Markdown("""<h1><center>My Chatbot</center></h1>""")

    chatbot = gr.Chatbot()

    state = gr.State()

    text = gr.Textbox(placeholder="Hello. Ask me a question.")

    submit = gr.Button("SEND")

    submit.click(my_chatbot, inputs=[text, state], outputs=[chatbot, state])

 

demo.launch(debug=True)