# visuals.py
# Handles all data visualization logic for MudabbirAI
import pandas as pd
import plotly.graph_objects as go
import config # Import config to access specific model names

def create_progress_chart(log_data):
    """
    Generates a Radar Chart comparing the initial draft scores vs. the final scores.
    """
    if not log_data or "trace" not in log_data:
        return None
    
    attempts = [step for step in log_data["trace"] if step["step_type"] == "attempt"]
    if not attempts:
        return None

    categories = ["Novelty", "Usefulness_Feasibility", "Flexibility", "Elaboration", "Cultural_Appropriateness"]
    
    fig = go.Figure()

    for i, attempt in enumerate(attempts):
        scores = attempt.get("scores", {})
        values = [scores.get(cat, 0) for cat in categories]
        
        # Close the loop for radar chart
        values += [values[0]]
        radar_categories = categories + [categories[0]]
        
        name = f"Initial Draft" if i == 0 else f"Improved Draft (Loop {i})"
        color = "red" if i == 0 else "green"
        
        fig.add_trace(go.Scatterpolar(
            r=values,
            theta=radar_categories,
            fill='toself',
            name=name,
            line_color=color
        ))

    fig.update_layout(
        polar=dict(
            radialaxis=dict(
                visible=True,
                range=[0, 5] 
            )),
        showlegend=True,
        title="Evolution of Solution Quality"
    )
    return fig

def create_calibration_table(log_data):
    """
    Generates a Pandas DataFrame showing the 'Audition Scores' for the calibration phase.
    """
    if not log_data or "trace" not in log_data:
        return None
    
    calibration_step = next((step for step in log_data["trace"] if step["step_type"] == "calibration"), None)
    if not calibration_step or "details" not in calibration_step:
        return None

    details = calibration_step["details"]
    data = []
    for item in details:
        role = item["role"]
        model_provider = item["llm"] # e.g., "Gemini", "Anthropic"
        score_data = item.get("score", {})
        
        score = 0
        if isinstance(score_data, dict):
             if role == "Plant": 
                 score = score_data.get("Novelty", {}).get("score", 0)
             elif role == "Implementer": 
                 score = score_data.get("Usefulness_Feasibility", {}).get("score", 0)
             elif role == "Monitor": 
                 score = score_data.get("Cultural_Appropriateness", {}).get("score", 0)
        
        # --- NEW: Retrieve specific model name from config for transparency ---
        specific_model_name = config.MODELS.get(model_provider, {}).get("default", "")
        if specific_model_name:
            # Format: "Anthropic\n(claude-3-5-haiku...)"
            display_name = f"{model_provider}\n({specific_model_name})"
        else:
            display_name = model_provider

        data.append({"Role": role, "Model": display_name, "Score": score})

    if not data: return None

    df = pd.DataFrame(data)
    pivot_df = df.pivot(index="Role", columns="Model", values="Score").reset_index()
    return pivot_df

def create_cost_summary(log_data):
    """
    Creates a Markdown summary of costs based on usage data found in the log.
    """
    if not log_data or "financial_report" not in log_data:
        return "### ⚠️ Financial Data Unavailable\n*Could not retrieve usage statistics.*"
    
    fin = log_data["financial_report"]
    total = fin.get("total_cost", 0.0)
    calib = fin.get("calibration_cost", 0.0)
    gen = fin.get("generation_cost", 0.0)
    
    # Calculate Total Tokens
    breakdown = fin.get("usage_breakdown", [])
    total_input = sum(u.get("input", 0) for u in breakdown)
    total_output = sum(u.get("output", 0) for u in breakdown)
    
    # Model Usage Breakdown
    models_used = {}
    for u in breakdown:
        m = u.get("model", "Unknown")
        models_used[m] = models_used.get(m, 0) + 1
    
    if not models_used:
        model_str = "None recorded"
    else:
        model_str = ", ".join([f"{k} ({v} calls)" for k,v in models_used.items()])

    md = f"""
    ### 💰 Financial Intelligence Report
    
    | **Category** | **Cost (USD)** | **Details** |
    | :--- | :--- | :--- |
    | **Total Investment** | **${total:.6f}** | **Total execution cost** |
    | Calibration Phase | ${calib:.6f} | Auditing models to pick the best team |
    | Solution Phase | ${gen:.6f} | Drafting, refining, and judging |
    
    ---
    **Operational Metrics:**
    * **Total Tokens:** {total_input + total_output:,} ({total_input:,} in / {total_output:,} out)
    * **Models Deployed:** {model_str}
    """
    return md