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
import logging

from .business_event_context import BusinessEventContext
from common.models.entitlements 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.document_version_file_size = kwargs.get('document_version_file_size')
        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 = []
        self.llm_metrics = {
            'total_tokens': 0,
            'prompt_tokens': 0,
            'completion_tokens': 0,
            'total_time': 0,
            'call_count': 0,
            'interaction_type': None
        }

    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}'")

    def update_llm_metrics(self, metrics: dict):
        self.llm_metrics['total_tokens'] += metrics['total_tokens']
        self.llm_metrics['prompt_tokens'] += metrics['prompt_tokens']
        self.llm_metrics['completion_tokens'] += metrics['completion_tokens']
        self.llm_metrics['total_time'] += metrics['time_elapsed']
        self.llm_metrics['call_count'] += 1
        self.llm_metrics['interaction_type'] = metrics['interaction_type']

    def reset_llm_metrics(self):
        self.llm_metrics['total_tokens'] = 0
        self.llm_metrics['prompt_tokens'] = 0
        self.llm_metrics['completion_tokens'] = 0
        self.llm_metrics['total_time'] = 0
        self.llm_metrics['call_count'] = 0
        self.llm_metrics['interaction_type'] = None

    @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:
            if self.llm_metrics['call_count'] > 0:
                self.log_final_metrics()
                self.reset_llm_metrics()
            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,
            'document_version_file_size': self.document_version_file_size,
            '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,
            document_version_file_size=self.document_version_file_size,
            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 log_llm_metrics(self, metrics: dict, level: str = 'info'):
        self.update_llm_metrics(metrics)
        message = "LLM Metrics"
        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,
            'document_version_file_size': self.document_version_file_size,
            'chat_session_id': self.chat_session_id,
            'interaction_id': self.interaction_id,
            'environment': self.environment,
            'llm_metrics_total_tokens': metrics['total_tokens'],
            'llm_metrics_prompt_tokens': metrics['prompt_tokens'],
            'llm_metrics_completion_tokens': metrics['completion_tokens'],
            'llm_metrics_total_time': metrics['time_elapsed'],
            'llm_interaction_type': metrics['interaction_type'],
        }
        # 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,
            document_version_file_size=self.document_version_file_size,
            chat_session_id=self.chat_session_id,
            interaction_id=self.interaction_id,
            environment=self.environment,
            llm_metrics_total_tokens=metrics['total_tokens'],
            llm_metrics_prompt_tokens=metrics['prompt_tokens'],
            llm_metrics_completion_tokens=metrics['completion_tokens'],
            llm_metrics_total_time=metrics['time_elapsed'],
            llm_interaction_type=metrics['interaction_type'],
            message=message
        )
        db.session.add(event_log)
        db.session.commit()

    def log_final_metrics(self, level: str = 'info'):
        logger = logging.getLogger('business_events')
        message = "Final LLM Metrics"
        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,
            'document_version_file_size': self.document_version_file_size,
            'chat_session_id': self.chat_session_id,
            'interaction_id': self.interaction_id,
            'environment': self.environment,
            'llm_metrics_total_tokens': self.llm_metrics['total_tokens'],
            'llm_metrics_prompt_tokens': self.llm_metrics['prompt_tokens'],
            'llm_metrics_completion_tokens': self.llm_metrics['completion_tokens'],
            'llm_metrics_total_time': self.llm_metrics['total_time'],
            'llm_metrics_call_count': self.llm_metrics['call_count'],
            'llm_interaction_type': self.llm_metrics['interaction_type'],
        }
        # 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,
            document_version_file_size=self.document_version_file_size,
            chat_session_id=self.chat_session_id,
            interaction_id=self.interaction_id,
            environment=self.environment,
            llm_metrics_total_tokens=self.llm_metrics['total_tokens'],
            llm_metrics_prompt_tokens=self.llm_metrics['prompt_tokens'],
            llm_metrics_completion_tokens=self.llm_metrics['completion_tokens'],
            llm_metrics_total_time=self.llm_metrics['total_time'],
            llm_metrics_call_count=self.llm_metrics['call_count'],
            llm_interaction_type=self.llm_metrics['interaction_type'],
            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):
        if self.llm_metrics['call_count'] > 0:
            self.log_final_metrics()
            self.reset_llm_metrics()
        self.log(f'Ending Trace for {self.event_type}')
        return BusinessEventContext(self).__exit__(exc_type, exc_val, exc_tb)