"""
BrainSteel Fin — Token Tracker
Logs LLM token usage per agent per call.
"""
import sqlite3, os
from datetime import datetime, timedelta

COSTS = {
    "google/gemma-4-31b-it:free": 0.0,
    "google/gemma-4-26b-a4b-it:free": 0.0,
    "deepseek/deepseek-v4-flash:free": 0.0,
    "anthropic/claude-sonnet-4": 3.0,
    "default": 0.0,
}

def log_tokens(agent: str, model: str, usage: dict):
    prompt_t = usage.get("prompt_tokens", 0)
    completion_t = usage.get("completion_tokens", 0)
    total_t = usage.get("total_tokens", prompt_t + completion_t)
    cost = COSTS.get(model, COSTS["default"]) * total_t / 1_000_000
    db_path = os.path.join(os.path.dirname(__file__), "..", "data", "executions.db")
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute(
        "INSERT INTO token_usage (timestamp, agent, model, prompt_tokens, completion_tokens, total_tokens, cost_usd) VALUES (?,?,?,?,?,?,?)",
        (datetime.now().isoformat(), agent, model, prompt_t, completion_t, total_t, cost)
    )
    conn.commit()
    conn.close()

def get_token_stats(days: int = 7) -> dict:
    db_path = os.path.join(os.path.dirname(__file__), "..", "data", "executions.db")
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    cutoff = (datetime.now() - timedelta(days=days)).isoformat()

    c.execute("""
        SELECT DATE(timestamp) as day, agent,
               SUM(prompt_tokens) as p, SUM(completion_tokens) as c, SUM(total_tokens) as t
        FROM token_usage
        WHERE timestamp >= ?
        GROUP BY day, agent
        ORDER BY day, agent
    """, (cutoff,))
    rows = c.fetchall()

    c.execute("""
        SELECT DATE(timestamp) as day,
               SUM(prompt_tokens), SUM(completion_tokens), SUM(total_tokens), SUM(cost_usd)
        FROM token_usage
        WHERE timestamp >= ?
        GROUP BY day ORDER BY day
    """, (cutoff,))
    daily = c.fetchall()

    c.execute("SELECT SUM(total_tokens), SUM(cost_usd), COUNT(*) FROM token_usage")
    total = c.fetchone()
    conn.close()

    agents = ["Brief", "Decio", "PaperT", "Audit"]
    days_set = sorted(set(r[0] for r in rows))
    chart_data = []
    for day in days_set:
        entry = {"date": day}
        for ag in agents:
            row = next((r for r in rows if r[0] == day and r[1] == ag), None)
            entry[ag] = row[4] if row else 0
        chart_data.append(entry)

    return {
        "total_all_time": total[0] or 0,
        "total_cost_usd": total[1] or 0.0,
        "total_calls": total[2] or 0,
        "daily_totals": [{"date": r[0], "prompt": r[1], "completion": r[2], "total": r[3], "cost": r[4]}
                         for r in daily],
        "chart_data": chart_data,
        "agents": agents,
    }