src/utils/session_tracker.py

"""
Session analytics tracker.

Tracks all API calls made during a session and provides cost/usage analytics.
"""

from typing import List, Dict, Optional
from datetime import datetime
import json


class SessionTracker:
    """Tracks API calls and provides analytics."""
    
    def __init__(self):
        self.calls: List[Dict] = []
        self.session_start = datetime.utcnow()
    
    def record_call(self, 
                   service: str,
                   action: str,
                   mode: str,  # 'mock' or 'real'
                   tokens_estimated: int,
                   cost_estimated: float,
                   model: str = 'claude-sonnet-4',
                   success: bool = True,
                   error: Optional[str] = None):
        """Record an API call."""
        call_record = {
            'timestamp': datetime.utcnow().isoformat(),
            'service': service,
            'action': action,
            'mode': mode,
            'tokens_estimated': tokens_estimated,
            'cost_estimated': cost_estimated,
            'model': model,
            'success': success,
            'error': error
        }
        self.calls.append(call_record)
    
    def get_summary(self) -> Dict:
        """Get session summary."""
        if not self.calls:
            return {
                'total_calls': 0,
                'session_duration': self._get_session_duration(),
                'message': 'No API calls recorded yet'
            }
        
        total_calls = len(self.calls)
        successful_calls = sum(1 for call in self.calls if call['success'])
        mock_calls = sum(1 for call in self.calls if call['mode'] == 'mock')
        real_calls = total_calls - mock_calls
        
        total_tokens = sum(call['tokens_estimated'] for call in self.calls)
        total_cost_saved = sum(call['cost_estimated'] for call in self.calls if call['mode'] == 'mock')
        total_cost_spent = sum(call['cost_estimated'] for call in self.calls if call['mode'] == 'real')
        
        # Breakdown by service
        by_service = {}
        for call in self.calls:
            service = call['service']
            if service not in by_service:
                by_service[service] = {
                    'calls': 0,
                    'tokens': 0,
                    'cost_saved': 0.0,
                    'cost_spent': 0.0
                }
            by_service[service]['calls'] += 1
            by_service[service]['tokens'] += call['tokens_estimated']
            if call['mode'] == 'mock':
                by_service[service]['cost_saved'] += call['cost_estimated']
            else:
                by_service[service]['cost_spent'] += call['cost_estimated']
        
        # Most called actions
        action_counts = {}
        for call in self.calls:
            key = f"{call['service']}.{call['action']}"
            action_counts[key] = action_counts.get(key, 0) + 1
        
        top_actions = sorted(action_counts.items(), key=lambda x: x[1], reverse=True)[:5]
        
        return {
            'session_start': self.session_start.isoformat(),
            'session_duration': self._get_session_duration(),
            'total_calls': total_calls,
            'successful_calls': successful_calls,
            'failed_calls': total_calls - successful_calls,
            'mock_calls': mock_calls,
            'real_calls': real_calls,
            'total_tokens_estimated': total_tokens,
            'total_cost_saved': f'${total_cost_saved:.4f}',
            'total_cost_spent': f'${total_cost_spent:.4f}',
            'savings_rate': f'{(mock_calls / total_calls * 100):.1f}%' if total_calls > 0 else '0%',
            'by_service': by_service,
            'top_actions': [{'action': action, 'count': count} for action, count in top_actions],
            'cost_breakdown': self._get_cost_breakdown()
        }
    
    def get_detailed_report(self) -> str:
        """Get detailed formatted report."""
        summary = self.get_summary()
        
        if summary.get('total_calls', 0) == 0:
            return "No API calls recorded yet. Start testing to see your savings!"
        
        report = []
        report.append("=" * 60)
        report.append("TOKEN ESTIMATOR - SESSION REPORT")
        report.append("=" * 60)
        report.append(f"\nSession Duration: {summary['session_duration']}")
        report.append(f"Total API Calls: {summary['total_calls']}")
        report.append(f"  ✓ Successful: {summary['successful_calls']}")
        report.append(f"  ✗ Failed: {summary['failed_calls']}")
        report.append(f"\nMode Breakdown:")
        report.append(f"  🎭 Mock Calls: {summary['mock_calls']}")
        report.append(f"  🌐 Real Calls: {summary['real_calls']}")
        
        report.append(f"\n💰 COST ANALYSIS")
        report.append(f"  Total Tokens: {summary['total_tokens_estimated']:,}")
        report.append(f"  Cost Saved (Mock): {summary['total_cost_saved']}")
        report.append(f"  Cost Spent (Real): {summary['total_cost_spent']}")
        report.append(f"  Savings Rate: {summary['savings_rate']}")
        
        if summary['by_service']:
            report.append(f"\n📊 BY SERVICE")
            for service, stats in summary['by_service'].items():
                report.append(f"\n  {service.upper()}:")
                report.append(f"    Calls: {stats['calls']}")
                report.append(f"    Tokens: {stats['tokens']:,}")
                report.append(f"    Saved: ${stats['cost_saved']:.4f}")
                if stats['cost_spent'] > 0:
                    report.append(f"    Spent: ${stats['cost_spent']:.4f}")
        
        if summary['top_actions']:
            report.append(f"\n🔥 MOST USED ACTIONS")
            for i, action_stat in enumerate(summary['top_actions'], 1):
                report.append(f"  {i}. {action_stat['action']}: {action_stat['count']} calls")
        
        report.append("\n" + "=" * 60)
        report.append("💡 TIP: Switch mock=False when ready for production!")
        report.append("=" * 60)
        
        return "\n".join(report)
    
    def _get_session_duration(self) -> str:
        """Get formatted session duration."""
        duration = datetime.utcnow() - self.session_start
        minutes = int(duration.total_seconds() / 60)
        seconds = int(duration.total_seconds() % 60)
        
        if minutes > 0:
            return f"{minutes}m {seconds}s"
        return f"{seconds}s"
    
    def _get_cost_breakdown(self) -> Dict:
        """Get cost breakdown by model."""
        breakdown = {}
        for call in self.calls:
            model = call['model']
            if model not in breakdown:
                breakdown[model] = {
                    'calls': 0,
                    'tokens': 0,
                    'cost': 0.0
                }
            breakdown[model]['calls'] += 1
            breakdown[model]['tokens'] += call['tokens_estimated']
            breakdown[model]['cost'] += call['cost_estimated']
        
        return breakdown
    
    def export_json(self) -> str:
        """Export session data as JSON."""
        return json.dumps({
            'summary': self.get_summary(),
            'calls': self.calls
        }, indent=2)
    
    def reset(self):
        """Reset session tracking."""
        self.calls = []
        self.session_start = datetime.utcnow()