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 = []
        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,
            '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 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,
            '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,
            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,
            '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,
            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)