
"""
Metrics Logger: Handles sub-agent usage and cost reporting.
"""

import json
import logging
from datetime import datetime
from typing import Dict
from pathlib import Path

from .config import METRICS_DIR

logger = logging.getLogger("geminis-metrics")

def log_metrics(prompt: str, output_data: Dict):
    """Log task execution metrics to docs/metrics/usage_report.md."""
    try:
        metrics_path = METRICS_DIR / "usage_report.md"
        if not metrics_path.parent.exists():
            metrics_path.parent.mkdir(parents=True)
        
        # Initialize if doesn't exist
        if not metrics_path.exists():
            metrics_path.write_text("# Sub-Agent Usage & Cost Report\n\n| Timestamp | Task | Model | Input Tokens | Output Tokens | Total Tokens | Tool Calls | Latency (ms) | Status |\n| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |\n")
        
        timestamp = datetime.now().isoformat()
        # Clean the task snippet for markdown table safety
        task_snippet = prompt[:100].replace("|", "\\|").replace("\n", " ")
        
        # Extract stats
        stats = output_data.get("stats", {})
        models = stats.get("models", {})
        
        # Aggregation
        model_name = "N/A"
        input_tokens = 0
        output_tokens = 0
        total_tokens = 0
        latency = 0
        
        if models:
            for name, data in models.items():
                model_name = name
                tokens = data.get("tokens", {})
                input_tokens += tokens.get("prompt", 0)
                output_tokens += tokens.get("candidates", 0)
                total_tokens += tokens.get("total", 0)
                latency += data.get("api", {}).get("totalLatencyMs", 0)
        else:
            # Check for generic stats if models list is empty
            input_tokens = stats.get("tokens", {}).get("prompt", 0)
            output_tokens = stats.get("tokens", {}).get("candidates", 0)
            total_tokens = stats.get("tokens", {}).get("total", 0)
        
        tools = stats.get("tools", {})
        tool_calls = tools.get("totalCalls", 0)
        status = output_data.get("status", "unknown")
        
        row = f"| {timestamp} | {task_snippet} | {model_name} | {input_tokens} | {output_tokens} | {total_tokens} | {tool_calls} | {latency} | {status} |\n"
        
        with open(metrics_path, "a") as f:
            f.write(row)
            
    except Exception as e:
        logger.warning(f"Failed to log metrics: {e}")
