"""
Lab 6.9: Cost Optimization -- SOLUTION

Profiles an agent session, implements cascade routing, calculates savings.
"""

MODEL_COSTS = {
    "gemini-flash": {"input": 0.15, "output": 0.60},
    "deepseek-v3": {"input": 0.27, "output": 1.10},
    "claude-sonnet": {"input": 3.00, "output": 15.00},
    "claude-opus": {"input": 15.00, "output": 75.00},
}

SESSION_LOG = [
    ("retrieve context", 8000, 200, "retrieval"),
    ("analyze data", 5000, 1500, "analysis"),
    ("make decision", 3000, 800, "decision"),
    ("format output", 2000, 400, "formatting"),
    ("verify result", 2500, 600, "verification"),
]

CASCADE_MAP = {
    "retrieval": "gemini-flash",
    "formatting": "gemini-flash",
    "verification": "deepseek-v3",
    "analysis": "claude-sonnet",
    "decision": "claude-opus",
}


def calculate_session_cost(session: list[tuple], model: str) -> float:
    costs = MODEL_COSTS[model]
    total = 0.0
    for step_name, in_tokens, out_tokens, _ in session:
        step_cost = (in_tokens * costs["input"] + out_tokens * costs["output"]) / 1_000_000
        total += step_cost
    return total


def calculate_cascade_cost(session: list[tuple]) -> float:
    total = 0.0
    for step_name, in_tokens, out_tokens, step_type in session:
        model = CASCADE_MAP[step_type]
        costs = MODEL_COSTS[model]
        step_cost = (in_tokens * costs["input"] + out_tokens * costs["output"]) / 1_000_000
        total += step_cost
        print(f"  {step_name:<25} {model:<16} in={in_tokens:>5} out={out_tokens:>5} -> ${step_cost:.6f}")
    return total


def find_optimizations(session: list[tuple], current_cost: float) -> list[dict]:
    optimizations = []
    for step_name, in_tokens, out_tokens, step_type in session:
        single_model = "claude-opus"
        cascade_model = CASCADE_MAP[step_type]
        
        sc = MODEL_COSTS[single_model]
        cc = MODEL_COSTS[cascade_model]
        
        single_cost = (in_tokens * sc["input"] + out_tokens * sc["output"]) / 1_000_000
        cascade_cost = (in_tokens * cc["input"] + out_tokens * cc["output"]) / 1_000_000
        savings = single_cost - cascade_cost
        
        if savings > 0.001:
            optimizations.append({
                "description": f"'{step_name}' from {single_model} to {cascade_model}",
                "savings": round(savings, 4),
                "savings_pct": round((savings / single_cost) * 100, 0)
            })
    
    optimizations.sort(key=lambda x: x["savings"], reverse=True)
    return optimizations


if __name__ == "__main__":
    print("=" * 55)
    print(" COST OPTIMIZATION ANALYSIS")
    print("=" * 55)
    
    opus_cost = calculate_session_cost(SESSION_LOG, "claude-opus")
    print(f"\nSingle model (all Opus): ${opus_cost:.4f}")
    
    print(f"\nCascade routing:")
    cascade_cost = calculate_cascade_cost(SESSION_LOG)
    
    savings_pct = ((opus_cost - cascade_cost) / opus_cost * 100)
    print(f"\nCascade total:    ${cascade_cost:.4f}")
    print(f"All Opus total:   ${opus_cost:.4f}")
    print(f"Savings:          {savings_pct:.0f}% (${opus_cost - cascade_cost:.4f})")
    
    optimizations = find_optimizations(SESSION_LOG, opus_cost)
    if optimizations:
        print(f"\nOptimization suggestions:")
        for opt in optimizations:
            print(f"  - {opt['description']}: save {opt['savings_pct']:.0f}% (${opt['savings']:.4f})")
    
    print(f"\n3x Rule estimate:")
    print(f"  Prototype cost:  ${cascade_cost:.4f}")
    print(f"  With retries:    ${cascade_cost * 1.5:.4f}")
    print(f"  Production:      ${cascade_cost * 3.0:.4f}")
    print(f"  At 1000 tasks:   ${cascade_cost * 3.0 * 1000:.2f}/month")