import os
import uuid
from contextlib import contextmanager
from datetime import datetime
from typing import Dict, Any, Optional
from datetime import datetime as dt, timezone as tz
from portkey_ai import Portkey, Config
import logging

from .business_event_context import BusinessEventContext
from common.models.monitoring import BusinessEventLog
from common.extensions import db


class BusinessEvent:
    # The BusinessEvent class itself is a context manager, but it doesn't use the @contextmanager decorator.
    # Instead, it defines __enter__ and __exit__ methods explicitly. This is because we're doing something a bit more
    # complex - we're interacting with the BusinessEventContext and the _business_event_stack.

    def __init__(self, event_type: str, tenant_id: int, **kwargs):
        self.event_type = event_type
        self.tenant_id = tenant_id
        self.trace_id = str(uuid.uuid4())
        self.span_id = None
        self.span_name = None
        self.parent_span_id = None
        self.document_version_id = kwargs.get('document_version_id')
        self.chat_session_id = kwargs.get('chat_session_id')
        self.interaction_id = kwargs.get('interaction_id')
        self.environment = os.environ.get("FLASK_ENV", "development")
        self.span_counter = 0
        self.spans = []

    def update_attribute(self, attribute: str, value: any):
        if hasattr(self, attribute):
            setattr(self, attribute, value)
        else:
            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{attribute}'")

    @contextmanager
    def create_span(self, span_name: str):
        # The create_span method is designed to be used as a context manager. We want to perform some actions when
        # entering the span (like setting the span ID and name) and some actions when exiting the span (like removing
        # these temporary attributes). The @contextmanager decorator allows us to write this method in a way that
        # clearly separates the "entry" and "exit" logic, with the yield statement in between.

        parent_span_id = self.span_id
        self.span_counter += 1
        new_span_id = str(uuid.uuid4())

        # Save the current span info
        self.spans.append((self.span_id, self.span_name, self.parent_span_id))

        # Set the new span info
        self.span_id = new_span_id
        self.span_name = span_name
        self.parent_span_id = parent_span_id

        self.log(f"Starting span {span_name}")

        try:
            yield
        finally:
            self.log(f"Ending span {span_name}")
            # Restore the previous span info
            if self.spans:
                self.span_id, self.span_name, self.parent_span_id = self.spans.pop()
            else:
                self.span_id = None
                self.span_name = None
                self.parent_span_id = None

    def log(self, message: str, level: str = 'info'):
        logger = logging.getLogger('business_events')
        log_data = {
            'event_type': self.event_type,
            'tenant_id': self.tenant_id,
            'trace_id': self.trace_id,
            'span_id': self.span_id,
            'span_name': self.span_name,
            'parent_span_id': self.parent_span_id,
            'document_version_id': self.document_version_id,
            'chat_session_id': self.chat_session_id,
            'interaction_id': self.interaction_id,
            'environment': self.environment
        }
        # log to Graylog
        getattr(logger, level)(message, extra=log_data)

        # Log to database
        event_log = BusinessEventLog(
            timestamp=dt.now(tz=tz.utc),
            event_type=self.event_type,
            tenant_id=self.tenant_id,
            trace_id=self.trace_id,
            span_id=self.span_id,
            span_name=self.span_name,
            parent_span_id=self.parent_span_id,
            document_version_id=self.document_version_id,
            chat_session_id=self.chat_session_id,
            interaction_id=self.interaction_id,
            environment=self.environment,
            message=message
        )
        db.session.add(event_log)
        db.session.commit()

    def __enter__(self):
        self.log(f'Starting Trace for {self.event_type}')
        return BusinessEventContext(self).__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.log(f'Ending Trace for {self.event_type}')
        return BusinessEventContext(self).__exit__(exc_type, exc_val, exc_tb)