import streamlit as st
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pandas as pd
import numpy as np
import json
import uuid
from datetime import datetime, timedelta
import time
from huggingface_hub import HfApi, login
from streamlit_option_menu import option_menu
import requests
import hashlib
import os

# Configure page
st.set_page_config(
    page_title="ML Tracker - Free W&B Alternative",
    page_icon="📊",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialize session state
if 'authenticated' not in st.session_state:
    st.session_state.authenticated = False
if 'user_token' not in st.session_state:
    st.session_state.user_token = None
if 'api_key' not in st.session_state:
    st.session_state.api_key = None
if 'experiments' not in st.session_state:
    st.session_state.experiments = {}
if 'current_experiment' not in st.session_state:
    st.session_state.current_experiment = None

# Custom CSS for better styling
st.markdown("""
<style>
    .main-header {
        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
        padding: 1rem;
        border-radius: 10px;
        margin-bottom: 1rem;
    }
    .metric-card {
        background: #f8f9fa;
        padding: 1rem;
        border-radius: 8px;
        border-left: 4px solid #667eea;
        margin-bottom: 1rem;
    }
    .api-key-box {
        background: #f1f3f4;
        padding: 1rem;
        border-radius: 8px;
        font-family: monospace;
        margin: 1rem 0;
    }
    .stButton > button {
        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
        color: white;
        border: none;
        border-radius: 6px;
        padding: 0.5rem 1rem;
        font-weight: 500;
    }
</style>
""", unsafe_allow_html=True)

def generate_api_key(user_token):
    """Generate a unique API key for the user"""
    return hashlib.sha256(f"{user_token}_{datetime.now().isoformat()}".encode()).hexdigest()[:32]

def authenticate_user():
    """Handle HuggingFace authentication"""
    st.markdown('<div class="main-header"><h1 style="color: white; margin: 0;">🤗 ML Tracker</h1><p style="color: white; margin: 0;">Free W&B Alternative on HuggingFace Spaces</p></div>', unsafe_allow_html=True)
    
    col1, col2, col3 = st.columns([1, 2, 1])
    
    with col2:
        st.markdown("### 🔐 Connect with HuggingFace")
        st.markdown("Enter your HuggingFace token to get started with experiment tracking!")
        
        hf_token = st.text_input(
            "HuggingFace Token",
            type="password",
            placeholder="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
            help="Get your token from https://huggingface.co/settings/tokens"
        )
        
        if st.button("🚀 Connect & Generate API Key", use_container_width=True):
            if hf_token:
                try:
                    # Verify token
                    api = HfApi(token=hf_token)
                    user_info = api.whoami()
                    
                    # Store authentication
                    st.session_state.authenticated = True
                    st.session_state.user_token = hf_token
                    st.session_state.api_key = generate_api_key(hf_token)
                    st.session_state.username = user_info['name']
                    
                    st.success(f"✅ Successfully connected as {user_info['name']}!")
                    time.sleep(1)
                    st.rerun()
                    
                except Exception as e:
                    st.error(f"❌ Authentication failed: {str(e)}")
            else:
                st.error("Please enter your HuggingFace token")

def show_api_key():
    """Display API key and usage instructions"""
    st.markdown("### 🔑 Your API Key")
    st.markdown(f'<div class="api-key-box"><strong>API Key:</strong> {st.session_state.api_key}</div>', unsafe_allow_html=True)
    
    st.markdown("### 📋 Usage Instructions")
    st.code(f"""
# Install the client
pip install requests

# Python usage example
import requests
import json

API_KEY = "{st.session_state.api_key}"
BASE_URL = "https://your-space-url.hf.space"

# Log metrics
def log_metrics(experiment_name, step, metrics):
    response = requests.post(
        f"{BASE_URL}/api/log",
        json={{
            "api_key": API_KEY,
            "experiment": experiment_name,
            "step": step,
            "metrics": metrics,
            "timestamp": time.time()
        }}
    )
    return response.json()

# Example usage
log_metrics("my_experiment", 1, {{
    "loss": 0.5,
    "accuracy": 0.85,
    "learning_rate": 0.001
}})
""", language="python")

def generate_sample_data():
    """Generate sample experiment data for demonstration"""
    experiments = {
        "cnn_image_classification": {
            "created_at": datetime.now() - timedelta(days=2),
            "metrics": [],
            "config": {
                "model": "ResNet50",
                "dataset": "CIFAR-10",
                "epochs": 100,
                "batch_size": 32,
                "learning_rate": 0.001
            }
        },
        "nlp_sentiment_analysis": {
            "created_at": datetime.now() - timedelta(days=1),
            "metrics": [],
            "config": {
                "model": "BERT",
                "dataset": "IMDB",
                "epochs": 50,
                "batch_size": 16,
                "learning_rate": 0.0001
            }
        }
    }
    
    # Generate sample metrics
    for exp_name, exp_data in experiments.items():
        metrics = []
        for step in range(1, 101):
            if exp_name == "cnn_image_classification":
                loss = 2.3 * np.exp(-step/20) + 0.1 + np.random.normal(0, 0.05)
                accuracy = 1 - 0.9 * np.exp(-step/15) + np.random.normal(0, 0.02)
                val_loss = loss + np.random.normal(0, 0.1)
                val_accuracy = accuracy - np.random.normal(0.05, 0.02)
                
                metrics.append({
                    "step": step,
                    "loss": max(0, loss),
                    "accuracy": max(0, min(1, accuracy)),
                    "val_loss": max(0, val_loss),
                    "val_accuracy": max(0, min(1, val_accuracy)),
                    "timestamp": (datetime.now() - timedelta(days=2) + timedelta(minutes=step*2)).isoformat()
                })
            else:
                loss = 1.8 * np.exp(-step/25) + 0.2 + np.random.normal(0, 0.03)
                f1_score = 1 - 0.7 * np.exp(-step/20) + np.random.normal(0, 0.02)
                precision = f1_score + np.random.normal(0, 0.02)
                recall = f1_score + np.random.normal(0, 0.02)
                
                metrics.append({
                    "step": step,
                    "loss": max(0, loss),
                    "f1_score": max(0, min(1, f1_score)),
                    "precision": max(0, min(1, precision)),
                    "recall": max(0, min(1, recall)),
                    "timestamp": (datetime.now() - timedelta(days=1) + timedelta(minutes=step*3)).isoformat()
                })
        
        exp_data["metrics"] = metrics
    
    return experiments

def create_metric_charts(experiment_data):
    """Create interactive charts for experiment metrics"""
    if not experiment_data["metrics"]:
        st.warning("No metrics data available for this experiment.")
        return
    
    df = pd.DataFrame(experiment_data["metrics"])
    
    # Get all numeric columns (excluding step and timestamp)
    numeric_cols = [col for col in df.columns if col not in ['step', 'timestamp'] and pd.api.types.is_numeric_dtype(df[col])]
    
    if not numeric_cols:
        st.warning("No numeric metrics found.")
        return
    
    # Create subplots
    n_metrics = len(numeric_cols)
    n_cols = 2
    n_rows = (n_metrics + n_cols - 1) // n_cols
    
    fig = make_subplots(
        rows=n_rows,
        cols=n_cols,
        subplot_titles=numeric_cols,
        vertical_spacing=0.1,
        horizontal_spacing=0.1
    )
    
    colors = px.colors.qualitative.Set3
    
    for i, metric in enumerate(numeric_cols):
        row = i // n_cols + 1
        col = i % n_cols + 1
        
        fig.add_trace(
            go.Scatter(
                x=df['step'],
                y=df[metric],
                mode='lines+markers',
                name=metric,
                line=dict(color=colors[i % len(colors)], width=2),
                marker=dict(size=4),
                hovertemplate=f"<b>{metric}</b><br>Step: %{{x}}<br>Value: %{{y:.4f}}<extra></extra>"
            ),
            row=row,
            col=col
        )
    
    fig.update_layout(
        height=400 * n_rows,
        showlegend=False,
        title_text="Experiment Metrics Over Time",
        title_x=0.5,
        font=dict(size=12)
    )
    
    fig.update_xaxes(title_text="Step")
    fig.update_yaxes(title_text="Value")
    
    st.plotly_chart(fig, use_container_width=True)

def show_experiment_dashboard():
    """Display the main experiment dashboard"""
    st.markdown('<div class="main-header"><h1 style="color: white; margin: 0;">📊 ML Experiment Dashboard</h1></div>', unsafe_allow_html=True)
    
    # Load sample data if no experiments exist
    if not st.session_state.experiments:
        st.session_state.experiments = generate_sample_data()
    
    # Sidebar for experiment selection
    with st.sidebar:
        st.markdown("### 🔬 Experiments")
        
        exp_names = list(st.session_state.experiments.keys())
        if exp_names:
            selected_exp = st.selectbox(
                "Select Experiment",
                exp_names,
                key="exp_selector"
            )
            st.session_state.current_experiment = selected_exp
        else:
            st.info("No experiments found. Start logging metrics to see them here!")
            return
        
        st.markdown("### 📈 Quick Stats")
        if st.session_state.current_experiment:
            exp_data = st.session_state.experiments[st.session_state.current_experiment]
            st.metric("Total Steps", len(exp_data["metrics"]))
            st.metric("Created", exp_data["created_at"].strftime("%Y-%m-%d"))
    
    # Main dashboard content
    if st.session_state.current_experiment:
        exp_data = st.session_state.experiments[st.session_state.current_experiment]
        
        # Experiment header
        col1, col2 = st.columns([3, 1])
        with col1:
            st.markdown(f"## {st.session_state.current_experiment}")
        with col2:
            if st.button("🔄 Refresh", use_container_width=True):
                st.rerun()
        
        # Configuration section
        with st.expander("⚙️ Configuration", expanded=False):
            config_df = pd.DataFrame(list(exp_data["config"].items()), columns=["Parameter", "Value"])
            st.dataframe(config_df, use_container_width=True)
        
        # Metrics overview
        if exp_data["metrics"]:
            latest_metrics = exp_data["metrics"][-1]
            
            st.markdown("### 📊 Latest Metrics")
            cols = st.columns(len([k for k in latest_metrics.keys() if k not in ['step', 'timestamp']]))
            
            for i, (key, value) in enumerate(latest_metrics.items()):
                if key not in ['step', 'timestamp']:
                    with cols[i]:
                        st.metric(key.replace('_', ' ').title(), f"{value:.4f}")
        
        # Charts section
        st.markdown("### 📈 Metrics Over Time")
        create_metric_charts(exp_data)
        
        # Raw data section
        with st.expander("📋 Raw Data", expanded=False):
            if exp_data["metrics"]:
                df = pd.DataFrame(exp_data["metrics"])
                st.dataframe(df, use_container_width=True)
            else:
                st.info("No metrics data available.")

def main():
    """Main application logic"""
    if not st.session_state.authenticated:
        authenticate_user()
    else:
        # Navigation menu
        selected = option_menu(
            menu_title=None,
            options=["Dashboard", "API Key", "Logout"],
            icons=["graph-up", "key", "box-arrow-right"],
            menu_icon="cast",
            default_index=0,
            orientation="horizontal",
            styles={
                "container": {"padding": "0!important", "background-color": "#fafafa"},
                "icon": {"color": "#667eea", "font-size": "18px"},
                "nav-link": {"font-size": "16px", "text-align": "center", "margin": "0px", "--hover-color": "#eee"},
                "nav-link-selected": {"background-color": "#667eea"},
            }
        )
        
        if selected == "Dashboard":
            show_experiment_dashboard()
        elif selected == "API Key":
            show_api_key()
        elif selected == "Logout":
            st.session_state.authenticated = False
            st.session_state.user_token = None
            st.session_state.api_key = None
            st.session_state.experiments = {}
            st.rerun()

if __name__ == "__main__":
    main()