import logging
import time
from typing import Any, Dict, Literal, Optional

from opentelemetry import context as context_api
from opentelemetry import trace
from opentelemetry.trace import SpanKind, Status, StatusCode
from opentelemetry.trace.propagation import set_span_in_context

from netra.config import Config

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class ATTRIBUTE:
    LLM_SYSTEM = "llm_system"
    MODEL = "model"
    PROMPT = "prompt"
    NEGATIVE_PROMPT = "negative_prompt"
    IMAGE_HEIGHT = "image_height"
    IMAGE_WIDTH = "image_width"
    CREDITS = "credits"
    TOTAL_COST = "total_cost"
    TOTAL_TOKENS = "total_tokens"
    PROMPT_TOKENS_COST = "prompt_tokens_cost"
    PROMPT_TOKENS = "prompt_tokens"
    COMPLETION_TOKENS_COST = "completion_tokens_cost"
    COMPLETION_TOKENS = "completion_tokens"
    CACHED_TOKENS_COST = "cached_tokens_cost"
    CACHED_TOKENS = "cached_tokens"
    STATUS = "status"
    DURATION_MS = "duration_ms"
    ERROR_MESSAGE = "error_message"


class Session:
    """
    Context manager for tracking observability data for external API calls.

    Usage:
        with combat.start_session("video_gen_task") as session:
            session.set_prompt("A cat playing piano").set_image_height("1024")

            # External API call
            result = external_api.generate_video(...)

            session.set_tokens("20").set_credits("30")
    """

    def __init__(self, name: str, attributes: Optional[Dict[str, str]] = None, module_name: str = "combat_sdk"):
        self.name = name
        self.attributes = attributes or {}
        self.start_time: Optional[float] = None
        self.end_time: Optional[float] = None
        self.status = "pending"
        self.error_message: Optional[str] = None
        self.module_name = module_name

        # OpenTelemetry span management
        self.tracer = trace.get_tracer(module_name)
        self.span: Optional[trace.Span] = None
        self.context_token: Optional[Any] = None

    def __enter__(self) -> "Session":
        """Start the session, begin time tracking, and create OpenTelemetry span."""
        self.start_time = time.time()

        # Create OpenTelemetry span
        self.span = self.tracer.start_span(name=self.name, kind=SpanKind.CLIENT, attributes=self.attributes)

        # Set span in context
        ctx = set_span_in_context(self.span)
        self.context_token = context_api.attach(ctx)

        logger.info(f"Started session: {self.name}")
        return self

    def __exit__(self, exc_type: Optional[type], exc_val: Optional[Exception], exc_tb: Any) -> Literal[False]:
        """End the session, calculate duration, handle errors, and close OpenTelemetry span."""
        self.end_time = time.time()
        duration_ms = (self.end_time - self.start_time) * 1000 if self.start_time is not None else None

        # Set duration
        if duration_ms is not None:
            self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.DURATION_MS}", str(round(duration_ms, 2)))

        # Handle status and errors
        if exc_type is None and self.status == "pending":
            self.status = "success"
            if self.span:
                self.span.set_status(Status(StatusCode.OK))
        elif exc_type is not None:
            self.status = "error"
            self.error_message = str(exc_val)
            self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.ERROR_MESSAGE}", self.error_message)
            if self.span:
                self.span.set_status(Status(StatusCode.ERROR, self.error_message))
                if exc_val is not None:
                    self.span.record_exception(exc_val)
            logger.error(f"Session {self.name} failed: {self.error_message}")

        self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.STATUS}", self.status)

        # Update span attributes with final values
        if self.span:
            for key, value in self.attributes.items():
                self.span.set_attribute(key, value)

        # End OpenTelemetry span and detach context
        if self.span:
            self.span.end()
        if self.context_token:
            context_api.detach(self.context_token)

        logger.info(
            f"Ended session: {self.name} (Status: {self.status}, Duration: {duration_ms:.2f}ms)"
            if duration_ms is not None
            else f"Ended session: {self.name} (Status: {self.status})"
        )

        # Don't suppress exceptions
        return False

    def set_attribute(self, key: str, value: str) -> "Session":
        """Set a single attribute and return self for method chaining."""
        self.attributes[key] = value
        # Also set on the span if it exists
        if self.span:
            self.span.set_attribute(key, value)
        return self

    def set_prompt(self, prompt: str) -> "Session":
        """Set the input prompt."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.PROMPT}", prompt)

    def set_negative_prompt(self, negative_prompt: str) -> "Session":
        """Set the negative prompt."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.NEGATIVE_PROMPT}", negative_prompt)

    def set_image_height(self, height: str) -> "Session":
        """Set the image height."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.IMAGE_HEIGHT}", height)

    def set_image_width(self, width: str) -> "Session":
        """Set the image width."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.IMAGE_WIDTH}", width)

    def set_total_tokens(self, tokens: str) -> "Session":
        """Set the number of tokens used."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.TOTAL_TOKENS}", tokens)

    def set_prompt_tokens_cost(self, cost: str) -> "Session":
        """Set the number of tokens used."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.PROMPT_TOKENS_COST}", cost)

    def set_prompt_tokens(self, tokens: str) -> "Session":
        """Set the number of tokens used."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.PROMPT_TOKENS}", tokens)

    def set_completion_tokens_cost(self, cost: str) -> "Session":
        """Set the number of tokens used."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.COMPLETION_TOKENS_COST}", cost)

    def set_completion_tokens(self, tokens: str) -> "Session":
        """Set the number of tokens used."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.COMPLETION_TOKENS}", tokens)

    def set_cached_tokens_cost(self, cost: str) -> "Session":
        """Set the number of tokens used."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.CACHED_TOKENS_COST}", cost)

    def set_cached_tokens(self, tokens: str) -> "Session":
        """Set the number of tokens used."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.CACHED_TOKENS}", tokens)

    def set_credits(self, credits: str) -> "Session":
        """Set the number of credits used."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.CREDITS}", credits)

    def set_total_cost(self, cost: str) -> "Session":
        """Set the cost of the operation."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.TOTAL_COST}", cost)

    def set_model(self, model: str) -> "Session":
        """Set the model used."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.MODEL}", model)

    def set_llm_system(self, system: str) -> "Session":
        """Set the LLM system used."""
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.LLM_SYSTEM}", system)

    def set_error(self, error_message: str) -> "Session":
        """Manually set an error message."""
        self.status = "error"
        self.error_message = error_message
        if self.span:
            self.span.set_status(Status(StatusCode.ERROR, error_message))
        return self.set_attribute(f"{Config.LIBRARY_NAME}.{ATTRIBUTE.ERROR_MESSAGE}", error_message)

    def set_success(self) -> "Session":
        """Manually mark the session as successful."""
        self.status = "success"
        if self.span:
            self.span.set_status(Status(StatusCode.OK))
        return self

    def add_event(self, name: str, attributes: Optional[Dict[str, str]] = None) -> "Session":
        """Add an event to the span."""
        if self.span:
            self.span.add_event(name, attributes or {})
        return self

    def get_current_span(self) -> Optional[trace.Span]:
        """Get the current OpenTelemetry span."""
        return self.span
