import gradio as gr
import pandas as pd
from bs4 import BeautifulSoup
from typing import Tuple

import src.few_shot_funcs as fsf
import src.scrape as scrape


def asin_to_pdp(asin_or_url: str) -> dict:
    if isinstance(asin_or_url, str) and len(asin_or_url) == 10:
        asin_url = 'https://www.amazon.com/dp/' + asin_or_url
    elif fsf.check_url_structure(asin_or_url):
        asin_url = asin_or_url
    else:
        raise gr.Error('You must provide a valid ASIN (10 char code) or URL')
        
    html = scrape.zyte_call(asin_url)
    asin_pdp = scrape.get_asin_pdp(BeautifulSoup(html, 'html.parser'))
    if not asin_pdp:
        raise gr.Error('Input URL not found (404)')
    elif not asin_pdp.get('title') or not asin_pdp.get('tech_data'):
        raise gr.Error("Couldn't fetch title or technical details from input URL")
    return asin_pdp


def generate_bullets(title: str, tech_data: gr.DataFrame) -> str:
    tech_str = fsf.format_tech_as_str(tech_data)
    
    feature_bullets = fsf.generate_data(title=title, tech_process=tech_str,
                                        few_shot_df=fsf.FS_DS, vector_db=fsf.DB)
    return feature_bullets


def asin_comparison(asin_or_url: str) -> Tuple[str, str]:
    asin_pdp = asin_to_pdp(asin_or_url)
    
    input_title = asin_pdp.get('title')
    tech_details = pd.DataFrame([(k, v) for k, v in asin_pdp.get('tech_data').items()])
    
    feature_bullets = generate_bullets(input_title, tech_details)
    comparison_bullets = "## Original Bullets\n- " + '\n- '.join(asin_pdp.get('feature_bullets'))
    return feature_bullets, comparison_bullets



demo = gr.Blocks()

with demo:
    gr.Markdown(
        """
    ### Generate Product Feature-Bullets!
    - Input an Amazon ASIN or product URL (amazon.com URL only), or you can input your own data - a short title and techincal details\n
    - Hit "Generate" and watch the AI engine do its magic 🪄 🪄
    """
    )
    with gr.Tab("Input ASIN"):
        with gr.Row():
            asin_input = gr.Textbox(type="text", label="Enter Product ASIN / URL")  
        
        with gr.Row():
            clear_button = gr.ClearButton(components=asin_input,
                                          value='Clear', variant='primary')
            gen_button = gr.Button("Generate!")

        with gr.Row():
            output = gr.Markdown()
            comparison_output = gr.Markdown()

            gen_button.click(fn=asin_comparison, inputs=asin_input, 
                             outputs=[output, comparison_output], show_progress=True)
            
    
    with gr.Tab("Input Data"):
        with gr.Row():
            title_input = gr.Textbox(type="text", label="Enter Product Title")
            tech_input = gr.Dataframe(label='Enter Technical Details',
                headers=["Feature", "Value"],
                datatype=["str", "str"],
                row_count=(2, "dynamic"),
                col_count=(2, "static"),
            )
        
        with gr.Row():
            clear_button = gr.ClearButton(components=[title_input, tech_input],
                                          value='Clear', variant='primary')
            gen_button = gr.Button("Generate!")
    
        output = gr.Markdown()

        gen_button.click(fn=generate_bullets, inputs=[title_input, tech_input], outputs=output,
                          show_progress=True)

demo.launch()