"""
Usage Tracker
Stores and retrieves interaction metadata in a local SQLite database.
"""
import sqlite3
import time
from pathlib import Path
from typing import Dict, Any, List

class UsageTracker:
    def __init__(self, data_dir: Path):
        self.data_dir = data_dir / "learning"
        self.data_dir.mkdir(parents=True, exist_ok=True)
        self.db_path = self.data_dir / "usage.db"
        self._init_db()

    def _init_db(self):
        with sqlite3.connect(self.db_path) as conn:
            c = conn.cursor()
            c.execute('''CREATE TABLE IF NOT EXISTS interactions (
                id INTEGER PRIMARY KEY,
                timestamp TEXT,
                model TEXT,
                strategy TEXT,
                prompt_tokens INT,
                response_tokens INT,
                latency_ms INT,
                status TEXT,
                session_id TEXT
            )''')
            conn.commit()

    def log(self, model: str, strategy: str, prompt_tokens: int, response_tokens: int, latency_ms: int, status: str, session_id: str):
        with sqlite3.connect(self.db_path) as conn:
            c = conn.cursor()
            c.execute('''INSERT INTO interactions
                         (timestamp, model, strategy, prompt_tokens, response_tokens, latency_ms, status, session_id)
                         VALUES (?, ?, ?, ?, ?, ?, ?, ?)''',
                      (time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), model, strategy, prompt_tokens, response_tokens, latency_ms, status, session_id))
            conn.commit()

    def get_stats(self, days: int = 30) -> Dict[str, Any]:
        with sqlite3.connect(self.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - days * 86400))
            c.execute('''SELECT COUNT(*), SUM(prompt_tokens + response_tokens), AVG(latency_ms),
                         SUM(CASE WHEN status='success' THEN 1 ELSE 0 END)*1.0/COUNT(*),
                         COUNT(DISTINCT session_id)
                         FROM interactions WHERE timestamp >= ?''', (cutoff,))
            row = c.fetchone()
            return {
                "total_interactions": row[0] or 0,
                "total_tokens": row[1] or 0,
                "avg_latency_ms": row[2] or 0.0,
                "success_rate": row[3] or 0.0,
                "unique_sessions": row[4] or 0
            }

    def top_models(self, days: int = 30, limit: int = 5) -> List[Dict[str, Any]]:
        with sqlite3.connect(self.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - days * 86400))
            c.execute('''SELECT model, COUNT(*) as cnt, AVG(latency_ms),
                         SUM(CASE WHEN status='success' THEN 1 ELSE 0 END)*1.0/COUNT(*)
                         FROM interactions WHERE timestamp >= ?
                         GROUP BY model ORDER BY cnt DESC LIMIT ?''', (cutoff, limit))
            return [{"model": r[0], "count": r[1], "avg_latency_ms": r[2], "success_rate": r[3]} for r in c.fetchall()]

    def top_strategies(self, days: int = 30, limit: int = 5) -> List[Dict[str, Any]]:
        with sqlite3.connect(self.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - days * 86400))
            c.execute('''SELECT strategy, COUNT(*) as cnt, AVG(latency_ms)
                         FROM interactions WHERE timestamp >= ?
                         GROUP BY strategy ORDER BY cnt DESC LIMIT ?''', (cutoff, limit))
            return [{"strategy": r[0], "count": r[1], "avg_latency_ms": r[2]} for r in c.fetchall()]

    def hourly_distribution(self, days: int = 30) -> Dict[int, int]:
        with sqlite3.connect(self.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - days * 86400))
            c.execute('''SELECT strftime('%H', timestamp), COUNT(*)
                         FROM interactions WHERE timestamp >= ?
                         GROUP BY strftime('%H', timestamp)''', (cutoff,))
            dist = {i: 0 for i in range(24)}
            for row in c.fetchall():
                if row[0] is not None:
                    dist[int(row[0])] = row[1]
            return dist

    def daily_trend(self, days: int = 30) -> List[Dict[str, Any]]:
        with sqlite3.connect(self.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - days * 86400))
            c.execute('''SELECT strftime('%Y-%m-%d', timestamp), COUNT(*), SUM(prompt_tokens + response_tokens), AVG(latency_ms)
                         FROM interactions WHERE timestamp >= ?
                         GROUP BY strftime('%Y-%m-%d', timestamp) ORDER BY strftime('%Y-%m-%d', timestamp)''', (cutoff,))
            return [{"date": r[0], "count": r[1], "tokens": r[2], "avg_latency_ms": r[3]} for r in c.fetchall()]

    def session_summary(self, session_id: str) -> Dict[str, Any]:
        with sqlite3.connect(self.db_path) as conn:
            c = conn.cursor()
            c.execute('''SELECT MIN(timestamp), MAX(timestamp), COUNT(*), SUM(prompt_tokens + response_tokens)
                         FROM interactions WHERE session_id = ?''', (session_id,))
            row = c.fetchone()
            if not row or row[2] == 0:
                return {}
            
            c.execute('SELECT DISTINCT model FROM interactions WHERE session_id = ?', (session_id,))
            models = [r[0] for r in c.fetchall()]
            
            c.execute('SELECT DISTINCT strategy FROM interactions WHERE session_id = ?', (session_id,))
            strategies = [r[0] for r in c.fetchall()]
            
            return {
                "session_id": session_id,
                "start_time": row[0],
                "end_time": row[1],
                "interactions": row[2],
                "models_used": models,
                "strategies_used": strategies,
                "total_tokens": row[3]
            }

    def get_sessions(self, days: int = 7, limit: int = 20) -> List[Dict[str, Any]]:
        with sqlite3.connect(self.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - days * 86400))
            c.execute('''SELECT session_id, MIN(timestamp), MAX(timestamp), COUNT(*)
                         FROM interactions WHERE timestamp >= ?
                         GROUP BY session_id ORDER BY MIN(timestamp) DESC LIMIT ?''', (cutoff, limit))
            sessions = []
            for r in c.fetchall():
                # rough duration calc
                try:
                    start = time.mktime(time.strptime(r[1], "%Y-%m-%dT%H:%M:%SZ"))
                    end = time.mktime(time.strptime(r[2], "%Y-%m-%dT%H:%M:%SZ"))
                    duration = (end - start) / 60.0
                except:
                    duration = 0.0
                sessions.append({
                    "session_id": r[0],
                    "start_time": r[1],
                    "interaction_count": r[3],
                    "duration_minutes": duration
                })
            return sessions

    def model_usage_over_time(self, model: str, days: int = 30) -> List[Dict[str, Any]]:
        with sqlite3.connect(self.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - days * 86400))
            c.execute('''SELECT strftime('%Y-%m-%d', timestamp), COUNT(*), AVG(latency_ms)
                         FROM interactions WHERE model = ? AND timestamp >= ?
                         GROUP BY strftime('%Y-%m-%d', timestamp) ORDER BY strftime('%Y-%m-%d', timestamp)''', (model, cutoff))
            return [{"date": r[0], "count": r[1], "avg_latency_ms": r[2]} for r in c.fetchall()]

    def export_csv(self, filepath: str, days: int = 30):
        import csv
        with sqlite3.connect(self.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - days * 86400))
            c.execute('SELECT * FROM interactions WHERE timestamp >= ? ORDER BY timestamp', (cutoff,))
            rows = c.fetchall()
            if rows:
                with open(filepath, 'w', newline='', encoding='utf-8') as f:
                    writer = csv.writer(f)
                    writer.writerow([d[0] for d in c.description])
                    writer.writerows(rows)

    def purge(self, older_than_days: int = 90) -> int:
        with sqlite3.connect(self.db_path) as conn:
            c = conn.cursor()
            cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - older_than_days * 86400))
            c.execute('DELETE FROM interactions WHERE timestamp < ?', (cutoff,))
            count = c.rowcount
            conn.commit()
            return count
