import streamlit as st
from langchain_openai import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from dotenv import load_dotenv
import json
import re


# Load environment variables
load_dotenv()

def getWebsiteLink(company_name):
    llm = OpenAI(temperature=0.0, model="gpt-3.5-turbo-instruct")

    prompt = PromptTemplate(
        input_variables=["company"],
        template="""
        Official website URL of {company} is

        RESPONSE:
        """,
    )

    chain = LLMChain(llm=llm, prompt=prompt)

    response = chain.run(company_name)

    urls = re.findall(r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+', response)

    if urls:  # If any URL is found
        # Take the first URL found (assuming the first URL is the most relevant)
        response_json = json.dumps({"website_link": urls[0]}, indent=4)
    else:
        # If no URL is found, return a "not found" message
        response_json = json.dumps({"website_link": "Not found"}, indent=4)

    return response_json


# UI Starts here
st.set_page_config(page_title="Website URL Generator", page_icon='🔗', layout='centered')

st.header("Website URL Generator")

company_name = st.text_area('Enter the company name', height=150)

submit = st.button("Generate official website")

if submit:
    company_link = getWebsiteLink(company_name)
    st.write("Generate official website:")
    st.code(company_link, language='json')