import axiom
from datetime import datetime,timedelta
import rfc3339
from dotenv import load_dotenv
import os



class AxiomLogger:

    def __init__(self):
        load_dotenv()
        api_token = os.getenv('AXIOM_API_KEY')
        org_id = os.getenv('AXIOM_ORG_ID')
        self.client = axiom.Client(api_token, org_id)
        self.dataset='sdxl-server'

    def get_timestamp_formatted(self):
        time = datetime.utcnow()
        time_formatted = rfc3339.format(time)
        return time_formatted
    
    def ingest_with_status(self, status, message, **kwargs):
        self.client.ingest_events(
            dataset=self.dataset,
            events=[{
                "_time": self.get_timestamp_formatted(),
                "status": status,
                "message": message,
                **kwargs
            }])

    def warning(self, message, **kwargs):
        self.ingest_with_status("WARNING", message, **kwargs)
    
    def error(self, message, **kwargs):
        self.ingest_with_status("ERROR", message, **kwargs)
    
    def info(self, message, **kwargs):
        self.ingest_with_status("INFO", message, **kwargs)

    def debug(self, message, **kwargs):
        self.ingest_with_status("DEBUG", message, **kwargs)
        
    def critical(self, message, **kwargs):
        self.ingest_with_status("CRITICAL", message, **kwargs)


def main():
    api_token = os.getenv('AXIOM_API_KEY')
    org_id = os.getenv('AXIOM_ORG_ID')
    client = axiom.Client(api_token, org_id)

    time = datetime.utcnow() - timedelta(hours=1)
    time_formatted = rfc3339.format(time)

    client.ingest_events(
        dataset="sdxl-server",
        events=[
            {"inference_test": "test1", "_time": time_formatted},
            {"inference_test": "test2", "_time": time_formatted},
        ])


if __name__=='__main__':
    main()