import gradio as gr
from scrape_amazon import AmazonScraper, AIInterface
import re


llms = ['meta-llama/Llama-2-70b-chat-hf', "mistralai/Mixtral-8x7B-Instruct-v0.1",
        "mistralai/Mistral-7B-Instruct-v0.1", "mistralai/Mistral-7B-Instruct-v0.2",
        "togethercomputer/LLaMA-2-7B-32K"]
scraper = AmazonScraper()
aii = AIInterface()

def replace_img_tag(url, html_text, img_dict):
    # Regular expression pattern to find <img> tags
    img_tag_pattern = r'<img src="(.*?)">'

    # Function to replace <img> tag
    def replacer(match, url=url):
        img_url = match.group(1)
        srcset = ', '.join(f'{url} {w}w' for url, (w, h) in img_dict.items())
        replacement = f'''<div style="display: flex; justify-content: center;">
    <figure>
     <a href="{url}"><img src="{img_url}" srcset="{srcset}" sizes="(max-width: 600px) 100vw, 600px" alt="Image from Amazon.com"></a>
     <figcaption style="text-align: center;">Image from Amazon.com</figcaption>
    </figure>
    </div>'''
        return replacement

    # Use re.sub() to replace <img> tags
    new_html_text = re.sub(img_tag_pattern, replacer, html_text)

    return new_html_text

def upload_to_wordpress(html_content):
    print("upload_to_wordpress() is called\n")
    print(html_content)
    pass

def write_article(url):
    # Your logic to fetch HTML content from the URL
    # Replace this with your actual implementation

    text = scraper.get_product_info_and_reviews(url)
    image = list(scraper.images)[0]
    # prompt_for_ai = ("Write a summary of the following product and an overview of people's experiences based on the "
    #                  "provided reviews of it as follows. Format it nicely and professionally in HTML:\n\n") + text

    # prompt_for_ai = "write a succinct summary article about this product. Format it nicely in HTML:\n\n" + text
    prompt_for_ai = ("write an article containing a short summary of the following product and an overview of people's "
                     "experiences based on the provided reviews of it as follows. Format it nicely in HTML:\n") + text

    # prompt_for_ai = (f"Write an HTML code that includes a professionally authored article summary of the following "
    #                  f"product and an overview of people's experiences based on the provided reviews of it as "
    #                  f"follows. Underneath the title add this tag <p class=\"has-small-font-size\">This page includes "
    #                  f"paid Amazon affiliate links</p> and Include a link to the product {url} at the very end. Also "
    #                  f"include this image {image} after the second paragraph. Format it nicely and professionally in "
    #                  f"HTML. :\n\n") + text
    html_content = aii.ask_ai(prompt_for_ai, model=llms[2])

    prompt_for_ai = (f"Take the following HTML code and slightly modify it. Underneath the title add this tag '<p "
                     f"class=\"has-small-font-size\">This page includes paid Amazon affiliate links</p>'. Include a "
                     f"link {url} to the product at the end. Also include this image {image} after the first or "
                     f"second paragraph. Return a nice and professional HTML code:\n") + html_content
    html_content = aii.ask_ai(prompt_for_ai, model=llms[2])
    html_content = replace_img_tag(url, html_content, scraper.images)
    print(html_content)
    return html_content


with gr.Blocks() as demo:
    gr.Markdown("Start typing below and then click **Run** to see the output.")

    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(placeholder="Amazon link:")
            btn_write_article = gr.Button("Write Article")
            btn_upload_to_wordpress = gr.Button("Upload To Wordpress")
        with gr.Column():
            html_output = gr.HTML()

    btn_write_article.click(write_article, inputs=[inp], outputs=html_output)
    btn_upload_to_wordpress.click(upload_to_wordpress, inputs=[html_output])


demo.launch(server_port=7373, share=True)