import requests
import json
import os 
from dotenv import load_dotenv
from datetime import datetime, timedelta


load_dotenv()

API_BASE_URL = "https://api.smith.langchain.com/api/v1"

class LangSmithAPI:
    _daily_count = 0
    _last_count_date = None

    def __init__(self):
        self.base_url = API_BASE_URL
        self.headers = {
            'X-API-Key': os.environ.get("LANGCHAIN_API_KEY")
        }

        self.getSessionId()

        if self._last_count_date != datetime.today().date():
            self._daily_count = 0
            self._last_count_date = datetime.today().date()

    def getSessionId(self):

        payload = {}

        response = requests.request("GET", f"{self.base_url}/sessions", headers=self.headers, data=payload)
        response_json = response.json()

        for session in response_json:
            session_id = session.get('id', None)
            name = session.get('name', 'Unnamed')
            if( name == os.environ.get("LANGCHAIN_PROJECT") ):
                self.session_id = session_id


    def getLastTraces(self, id_request: str = ''):

        if not hasattr(self, 'session_id') or not self.session_id:
            return 0
        
        payload = {
            "session": [self.session_id],
            "run_type":"chain",
            "select": [
                "id","trace_id","name","run_type","extra","start_time","inputs",
                "total_tokens","prompt_tokens","completion_tokens",
                "total_cost","prompt_cost","completion_cost","first_token_time"
            ],
            "order": "desc",
            "limit": 100
        }

        # Make request
        response = requests.request("POST", f"{self.base_url}/runs/query", headers=self.headers, data=json.dumps(payload))
        response_json = response.json()
           
        # filters for the traces that have the same id_request
        # filtered_traces = [trace for trace in response_json['runs'] if 'extra' in trace and 'metadata' in trace['extra'] and 'id_request' in trace['extra']['metadata'] and trace['extra']['metadata']['id_request'] == id_request]
        filtered_traces = [trace for trace in response_json['runs'] if trace.get('parent_run_ids') is not None and trace['parent_run_ids'] == []]
        
        return filtered_traces
        
    
    def getDailyMessages(self, id_request: str = ""):
        trace = self.getLastTraces(id_request)
        
        if trace == 0:
            return 0
        
        current_date = datetime.today().date()
        # current_date = current_date - timedelta(days=1) # Yesterday
        
        # Reset counter if it's a new day
        if self._last_count_date != current_date:
            self._daily_count = 0
            self._last_count_date = current_date
        
        # Count only new messages from the current day
        for t in trace:
            start_time = datetime.fromisoformat(t['start_time'])
            if start_time.date() == current_date:
                self._daily_count += 1

        return self._daily_count