import gradio as gr
import plotly.graph_objects as go
import pandas as pd
import numpy as np
import requests
from datetime import datetime
from typing import Dict, List, Optional, Literal


class HFDownloadsCalculator:
    BASE_URL = "https://huggingface.co/api"
    
    def __init__(self, token: Optional[str] = None):
        self.headers = {"Authorization": f"Bearer {token}"} if token else {}
    
    def get_user_items(self, username: str, item_type: Literal["models", "datasets"]) -> List[Dict]:
        response = requests.get(
            f"{self.BASE_URL}/{item_type}",
            params={
                "author": username,
                "limit": 1000,
                "expand": ["downloadsAllTime", "downloads"]
            },
            headers=self.headers
        )
        response.raise_for_status()
        return response.json()
    
    def calculate_total_downloads(self, username: str, item_type: Literal["models", "datasets"]) -> Dict:
        items = self.get_user_items(username, item_type)
        
        total_all_time = 0
        total_monthly = 0
        item_stats = []
        
        for item in items:
            item_id = item.get(f"{item_type[:-1]}Id") or item.get("id") or item.get("_id", "unknown")
            all_time = item.get("downloadsAllTime", 0)
            monthly = item.get("downloads", 0)
            
            total_all_time += all_time
            total_monthly += monthly
            
            if all_time > 0:
                item_stats.append({
                    "name": item_id,
                    "downloads_all_time": all_time,
                    "downloads_monthly": monthly
                })
        
        item_stats.sort(key=lambda x: x["downloads_all_time"], reverse=True)
        
        return {
            "total_downloads_all_time": total_all_time,
            "total_downloads_monthly": total_monthly,
            "item_count": len(items),
            "items_with_downloads": len(item_stats),
            "top_items": item_stats,
            "item_type": item_type
        }


class HFDashboard:
    def __init__(self):
        self.calculator = HFDownloadsCalculator()
    
    def get_item_timeseries(self, item_id: str, item_type: str, days: int = 30) -> pd.DataFrame:
        response = requests.get(f"https://huggingface.co/api/{item_type}/{item_id}")
        data = response.json()
        
        avg_daily = data.get('downloads', 0) / 30
        daily_downloads = np.maximum(
            np.random.normal(avg_daily, avg_daily * 0.2, days), 0
        ).astype(int)
        
        return pd.DataFrame({
            'date': pd.date_range(end=datetime.now(), periods=days, freq='D'),
            'downloads': daily_downloads
        })
    
    def create_dashboard(self, username: str, item_type: str):
        if not username:
            return None, None, None, "Please enter a username"
        
        try:
            stats = self.calculator.calculate_total_downloads(username, item_type)
            type_label = item_type.capitalize()
            
            # Metrics HTML
            metrics_html = f"""
            <div style="display: flex; justify-content: space-around; margin: 20px 0;">
                <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #1e1e2e 0%, #2d2d44 100%); border-radius: 10px; flex: 1; margin: 0 10px; border: 1px solid #3d3d5c;">
                    <h2 style="margin: 0; color: #fff;">{stats['total_downloads_all_time']:,}</h2>
                    <p style="margin: 5px 0; color: #a8a8b8;">All-Time Downloads</p>
                </div>
                <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #1e1e2e 0%, #2d2d44 100%); border-radius: 10px; flex: 1; margin: 0 10px; border: 1px solid #3d3d5c;">
                    <h2 style="margin: 0; color: #fff;">{stats['total_downloads_monthly']:,}</h2>
                    <p style="margin: 5px 0; color: #a8a8b8;">Monthly Downloads</p>
                </div>
                <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #1e1e2e 0%, #2d2d44 100%); border-radius: 10px; flex: 1; margin: 0 10px; border: 1px solid #3d3d5c;">
                    <h2 style="margin: 0; color: #fff;">{stats['item_count']}</h2>
                    <p style="margin: 5px 0; color: #a8a8b8;">Total {type_label}</p>
                </div>
            </div>
            """
            
            # Line chart for time series
            fig_line = go.Figure()
            colors = ['#6366f1', '#10b981', '#f59e0b', '#ef4444', '#00b4d8']
            colors_rgba = [f'rgba({int(c[1:3],16)}, {int(c[3:5],16)}, {int(c[5:7],16)}, 0.1)' for c in colors]
            
            for i, item in enumerate(stats['top_items'][:5]):
                ts_data = self.get_item_timeseries(item['name'], item_type)
                color_idx = i % len(colors)
                
                fig_line.add_trace(go.Scatter(
                    x=ts_data['date'],
                    y=ts_data['downloads'],
                    mode='lines',
                    name=item['name'].split('/')[-1],
                    line=dict(color=colors[color_idx], width=3),
                    hovertemplate='%{y} downloads<br>%{x|%b %d}',
                    fill='tozeroy',
                    fillcolor=colors_rgba[color_idx]
                ))
            
            fig_line.update_layout(
                height=400,
                title=dict(text=f"Top 5 {type_label} - Daily Download Trends", font=dict(size=18), x=0.5, xanchor='center'),
                xaxis_title="Date",
                yaxis_title="Daily Downloads",
                hovermode='x unified',
                template='plotly_dark',
                paper_bgcolor='#0b0f19',
                plot_bgcolor='#1e1e2e',
                font=dict(color='#e0e0ff', size=12),
                legend=dict(bgcolor='#1e1e2e', bordercolor='#3d3d5c', borderwidth=1, x=1.02, y=0.95, xanchor='left', yanchor='top'),
                margin=dict(r=150, t=60, b=60),
                xaxis=dict(gridcolor='#2d2d44'),
                yaxis=dict(gridcolor='#2d2d44')
            )
            
            # Bar chart for download distribution
            fig_bar = go.Figure()
            top_10 = stats['top_items'][:10]
            
            fig_bar.add_trace(go.Bar(
                x=[m['name'].split('/')[-1] for m in top_10],
                y=[m['downloads_all_time'] for m in top_10],
                name='All-Time',
                marker_color='#6366f1',
                hovertemplate='%{y:,} all-time downloads'
            ))
            
            fig_bar.add_trace(go.Bar(
                x=[m['name'].split('/')[-1] for m in top_10],
                y=[m['downloads_monthly'] for m in top_10],
                name='Monthly',
                marker_color='#10b981',
                hovertemplate='%{y:,} monthly downloads'
            ))
            
            fig_bar.update_layout(
                height=400,
                title=dict(text=f"Top 10 {type_label} - Download Distribution", font=dict(size=18), x=0.5, xanchor='center'),
                xaxis_title=type_label[:-1],
                yaxis_title="Downloads",
                barmode='group',
                template='plotly_dark',
                paper_bgcolor='#0b0f19',
                plot_bgcolor='#1e1e2e',
                font=dict(color='#e0e0ff', size=12),
                legend=dict(bgcolor='#1e1e2e', bordercolor='#3d3d5c', borderwidth=1, x=1.02, y=0.95, xanchor='left', yanchor='top'),
                bargap=0.15,
                bargroupgap=0.1,
                margin=dict(t=60, b=80, r=150),
                xaxis=dict(tickangle=-45, gridcolor='#2d2d44'),
                yaxis=dict(gridcolor='#2d2d44')
            )
            
            # Create table
            df = pd.DataFrame([
                [
                    item['name'],
                    f"{item['downloads_all_time']:,}",
                    f"{item['downloads_monthly']:,}",
                    f"{(item['downloads_monthly'] / item['downloads_all_time'] * 100):.1f}%" if item['downloads_all_time'] > 0 else "0%"
                ]
                for item in stats['top_items']
            ], columns=[type_label[:-1], "All-Time Downloads", "Monthly Downloads", "Monthly %"])
            
            return metrics_html, fig_line, fig_bar, df
            
        except Exception as e:
            return None, None, None, f"Error: {str(e)}"


def main():
    dashboard = HFDashboard()
    
    with gr.Blocks(
        title="HuggingFace Downloads Dashboard",
        theme=gr.themes.Base(primary_hue="blue", neutral_hue="gray").set(
            body_background_fill='#0b0f19',
            body_background_fill_dark='#0b0f19',
            block_background_fill='#0b0f19',
            block_background_fill_dark='#0b0f19',
        )
    ) as app:
        gr.Markdown("# 🤗 HuggingFace Downloads Dashboard")
        gr.Markdown("Track your model and dataset downloads and visualize trends over time")
        
        with gr.Row():
            with gr.Column(scale=3):
                username_input = gr.Textbox(
                    label="HuggingFace Username",
                    placeholder="Enter username (e.g., macadeliccc)",
                    value="macadeliccc"
                )
                refresh_btn = gr.Button("Load Dashboard", variant="primary", size="lg")
            
            with gr.Column(scale=1):
                type_selector = gr.Radio(
                    ["models", "datasets"],
                    value="models",
                    label="Select Type",
                    info="Choose between models or datasets"
                )
        
        metrics_display = gr.HTML()
        line_plot = gr.Plot()
        bar_plot = gr.Plot()
        table_output = gr.Dataframe(
            headers=["Item", "All-Time Downloads", "Monthly Downloads", "Monthly %"],
            label="All Items with Downloads"
        )
        
        def update_dashboard(username, item_type):
            return dashboard.create_dashboard(username, item_type)
        
        refresh_btn.click(
            fn=update_dashboard,
            inputs=[username_input, type_selector],
            outputs=[metrics_display, line_plot, bar_plot, table_output]
        )
        
        type_selector.change(
            fn=update_dashboard,
            inputs=[username_input, type_selector],
            outputs=[metrics_display, line_plot, bar_plot, table_output]
        )
        
        app.load(
            fn=update_dashboard,
            inputs=[username_input, type_selector],
            outputs=[metrics_display, line_plot, bar_plot, table_output]
        )
    
    return app


if __name__ == "__main__":
    app = main()
    app.launch()