"""
Lab 5.9: Set Up Agent Observability

Trace every tool call + LLM completion to a local SQLite database.
"""

import sqlite3
import json
from datetime import datetime, timezone
from typing import Any


# TODO 1: Define the database schema
DB_SCHEMA = """
CREATE TABLE IF NOT EXISTS tool_calls (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    session_id TEXT NOT NULL,
    turn_number INTEGER NOT NULL,
    tool_name TEXT NOT NULL,
    tool_params TEXT NOT NULL,
    tool_result TEXT,
    duration_ms INTEGER,
    timestamp TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS llm_completions (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    session_id TEXT NOT NULL,
    turn_number INTEGER NOT NULL,
    input_tokens INTEGER,
    output_tokens INTEGER,
    cached_tokens INTEGER,
    cost_cents REAL,
    timestamp TEXT NOT NULL
);

-- TODO: Add a sessions table with total cost, total tokens, status
"""


class ObservabilityTracker:
    """Tracks agent observability data to SQLite."""
    
    def __init__(self, db_path: str = "agent_observability.db"):
        self.db_path = db_path
        self.session_id = f"session_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}"
        self._init_db()
    
    def _init_db(self):
        """Initialize database with schema."""
        # TODO: Connect to SQLite and create tables
        pass
    
    def log_tool_call(self, turn: int, tool_name: str, params: dict, 
                      result: str, duration_ms: int):
        """Log a tool call."""
        # TODO: Insert tool call record
        pass
    
    def log_llm_completion(self, turn: int, input_tokens: int, 
                           output_tokens: int, cost_cents: float):
        """Log an LLM completion with token counts."""
        # TODO: Insert LLM completion record
        pass
    
    def get_session_summary(self) -> dict:
        """Get total cost, tool calls, tokens for this session."""
        # TODO: Query and return summary
        pass


# TODO: Decorate your agent with observability
def with_observability(agent_func):
    """Decorator that wraps an agent with observability tracking."""
    def wrapper(prompt: str):
        tracker = ObservabilityTracker()
        # TODO: Wrap the agent execution with tracking
        result = agent_func(prompt)
        return result
    return wrapper


if __name__ == "__main__":
    # Test the tracker
    tracker = ObservabilityTracker()
    tracker.log_tool_call(1, "read_file", {"path": "test.txt"}, "file contents", 150)
    tracker.log_llm_completion(1, 500, 200, 0.003)
    
    summary = tracker.get_session_summary()
    print(f"Session: {tracker.session_id}")
    print(f"Summary: {json.dumps(summary, indent=2)}")