import streamlit as st
from mock_data.sample_data import generate_discovered_models, estimate_inference_cost

def get_gpu_tier_justification(model_size_gb, gpu_tier):
    """
    Provide justification for GPU tier selection based on model size.

    Args:
        model_size_gb: Model size in GB
        gpu_tier: Selected GPU tier (T4, V100, A100, etc.)

    Returns:
        Justification string explaining the GPU tier choice
    """
    if model_size_gb <= 4:
        return f"""💡 <b>{gpu_tier} selected:</b> This {model_size_gb}GB model fits comfortably in 16GB memory.
        The T4 GPU (${'{:.2f}'.format(0.35)}/hr) provides cost-efficient inference with 200 QPS throughput,
        ideal for production deployment of compact models."""
    elif model_size_gb <= 10:
        return f"""💡 <b>{gpu_tier} selected:</b> This {model_size_gb}GB model requires 32GB GPU memory.
        The V100 (${'{:.2f}'.format(2.50)}/hr) offers 120 QPS throughput with balanced performance and cost,
        optimal for medium-sized models with moderate inference workloads."""
    elif model_size_gb <= 20:
        return f"""💡 <b>{gpu_tier} selected:</b> This {model_size_gb}GB model needs 40GB+ GPU memory.
        The A100 (${'{:.2f}'.format(4.10)}/hr) delivers 80 QPS with superior compute performance,
        recommended for large models requiring high throughput and low latency."""
    else:
        num_gpus = int(gpu_tier.replace("A100x", "")) if "x" in gpu_tier else 1
        total_cost = 4.10 * num_gpus
        memory_per_gpu = model_size_gb / num_gpus
        return f"""💡 <b>{gpu_tier} selected:</b> This {model_size_gb}GB model exceeds single-GPU capacity.
        Using {num_gpus}× A100 GPUs (${'{:.2f}'.format(total_cost)}/hr total) with ~{memory_per_gpu:.1f}GB per GPU
        enables tensor parallelism for distributed inference of very large models."""

def render_model_cards():
    """Render the discovered models as cards with export buttons."""

    models = generate_discovered_models()

    # Display 6 models in 2 rows of 3
    for row in range(2):
        cols = st.columns(3)
        for col_idx in range(3):
            idx = row * 3 + col_idx
            if idx < len(models):
                model = models[idx]
                with cols[col_idx]:
                    # Create a card-like container with border
                    with st.container(border=True):
                        st.markdown(f"#### {model['name']}")
                        st.caption(f"Quantization: {model['quantization']}")

                        # Calculate cost metrics
                        cost_info = estimate_inference_cost(model['size_gb'], model['quantization'])
                        baseline_cost = estimate_inference_cost(32.5, "FP32")["cost_per_1M"]
                        cost_savings_pct = ((baseline_cost - cost_info['cost_per_1M']) / baseline_cost) * 100

                        # Model metrics in a more compact layout
                        col1, col2 = st.columns(2)
                        with col1:
                            st.metric(label="Params", value=model['params'])
                            st.metric(label="Size", value=f"{model['size_gb']} GB")
                            st.metric(label="Accuracy", value=f"{model['accuracy']}%")
                        with col2:
                            st.metric(label="Cost/1M", value=f"${cost_info['cost_per_1M']:.2f}",
                                     delta=f"-{cost_savings_pct:.0f}%", delta_color="inverse")
                            st.metric(label="Throughput", value=f"{int(cost_info['throughput_qps'])} QPS")
                            st.metric(label="GPU Tier", value=cost_info['gpu_tier'])

                        # GPU tier justification
                        gpu_justification = get_gpu_tier_justification(model['size_gb'], cost_info['gpu_tier'])
                        st.markdown(f"""
                        <div style='background: rgba(99, 102, 241, 0.1); padding: 0.75rem; border-radius: 8px;
                                    border-left: 3px solid #6366f1; margin: 0.75rem 0; line-height: 1.5;'>
                            <small style='color: rgba(255,255,255,0.85); font-size: 0.85rem;'>{gpu_justification}</small>
                        </div>
                        """, unsafe_allow_html=True)

                        # Export button
                        if st.button(
                            "📥 Export Model",
                            key=f"export_{model['name']}",
                            use_container_width=True,
                            type="primary"
                        ):
                            # Simulate export by creating a download button
                            mock_model_data = f"# {model['name']}\n\n"
                            mock_model_data += f"**Parameters:** {model['params']}\n"
                            mock_model_data += f"**Accuracy:** {model['accuracy']}%\n"
                            mock_model_data += f"**Size:** {model['size_gb']} GB\n"
                            mock_model_data += f"**Quantization:** {model['quantization']}\n"
                            mock_model_data += f"\n---\n"
                            mock_model_data += "This is a mock model file for demonstration purposes.\n"

                            st.download_button(
                                label=f"⬇️ Download {model['name']}",
                                data=mock_model_data,
                                file_name=f"{model['name'].lower().replace(' ', '_')}.txt",
                                mime="text/plain",
                                key=f"download_{model['name']}",
                                use_container_width=True
                            )