import sqlite3
import threading
import time

class AgentMeter:
    def __init__(self, db_path="metering.db"):
        self.db_path = db_path
        self.lock = threading.Lock()
        self._init_db()

    def _init_db(self):
        with self.lock:
            conn = sqlite3.connect(self.db_path)
            try:
                with conn:
                    cursor = conn.cursor()
                    cursor.execute('''
                        CREATE TABLE IF NOT EXISTS usage (
                            id INTEGER PRIMARY KEY AUTOINCREMENT,
                            agent_id TEXT NOT NULL,
                            timestamp REAL NOT NULL,
                            time_ms REAL NOT NULL,
                            tokens INTEGER NOT NULL
                        )
                    ''')
            finally:
                conn.close()

    def track_usage(self, agent_id: str, time_ms: float, tokens: int):
        timestamp = time.time()
        with self.lock:
            conn = sqlite3.connect(self.db_path)
            try:
                with conn:
                    cursor = conn.cursor()
                    cursor.execute('''
                        INSERT INTO usage (agent_id, timestamp, time_ms, tokens)
                        VALUES (?, ?, ?, ?)
                    ''', (agent_id, timestamp, time_ms, tokens))
            finally:
                conn.close()

    def get_usage(self, agent_id: str = None):
        with self.lock:
            conn = sqlite3.connect(self.db_path)
            try:
                cursor = conn.cursor()
                if agent_id:
                    cursor.execute('SELECT * FROM usage WHERE agent_id = ?', (agent_id,))
                else:
                    cursor.execute('SELECT * FROM usage')
                columns = [col[0] for col in cursor.description]
                return [dict(zip(columns, row)) for row in cursor.fetchall()]
            finally:
                conn.close()