import gradio as gr
import requests
import pandas as pd
import plotly.express as px

class DownloadsTracker:
    def __init__(self):
        self.all_packages = {}
        self.current_packages = []
    
    def fetch_if_needed(self, packages):
        errors = []
        for pkg in packages:
            if pkg not in self.all_packages:
                try:
                    response = requests.get(f"https://pypistats.org/api/packages/{pkg}/overall")
                    if response.status_code != 200:
                        errors.append(f"Package not found: {pkg}")
                        continue
                    
                    data = response.json()["data"]
                    df = pd.DataFrame([
                        {"date": d["date"], "downloads": d["downloads"]} 
                        for d in data if d["category"] == "without_mirrors"
                    ])
                    df["date"] = pd.to_datetime(df["date"])
                    df = df.sort_values("date")
                    df["cumulative_downloads"] = df["downloads"].cumsum()
                    self.all_packages[pkg] = df
                except:
                    errors.append(f"Error fetching {pkg}")
        if len(errors) > 0:
            return "\n".join(errors)
        return None

    def render(self, package_list, use_log_scale):
        package_list = [p.strip() for p in package_list.split(",") if p.strip()]
        self.current_packages = package_list
        errors = self.fetch_if_needed(package_list)
        return self.plot(use_log_scale), errors, gr.update(visible=errors is not None)
    
    def plot(self, use_log_scale):
        fig = px.line(title="Cumulative Package Downloads")
        for pkg in self.current_packages:
            if pkg in self.all_packages:
                df = self.all_packages[pkg]
                fig.add_scatter(x=df["date"], y=df["cumulative_downloads"], name=pkg)
        
        if use_log_scale:
            fig.update_yaxes(type="log")
        else:
            fig.update_yaxes(type="linear")
        return fig

tracker = DownloadsTracker()

css = """
#textbox_id textarea {color: red}
#textbox_id span {background-color: red}
"""

with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    packages = gr.Textbox("transformers", label="Package names (comma-separated)")
    log_scale = gr.Checkbox(label="Use logarithmic scale on Y-axis", value=False)
    error_box = gr.Textbox(label="Errors:", interactive=False, visible=False, elem_id="textbox_id")
    render_btn = gr.Button("Render")
    plot = gr.Plot()
    render_btn.click(tracker.render, inputs=[packages, log_scale], outputs=[plot, error_box, error_box])
    log_scale.change(tracker.plot, inputs=[log_scale], outputs=[plot])

    demo.load(tracker.render, inputs=[packages, log_scale], outputs=[plot, error_box, error_box])

demo.launch()