import json
import logging
from functools import cached_property
from typing import Any

from ..utils import format_hex_id
from .spec import OtelGenAISpanSpecification
from .types import (
    ERROR,
    INFO,
    OPTIONAL,
    RECOMMENDED,
    REQUIRED,
    WARNING,
    AttributeMetadata,
    AttributeSeverity,
    AttributeSpec,
    AttributeValidationDetail,
    AttributeValidationDetails,
)
from .utils import is_valid_type


class OtelGenAISpecCheck:
    def __init__(self, raw_span: Any):
        """Initialize the specification checker with a raw span.

        Args:
            raw_span: A span object (dict or JSON string) to validate against the
                OpenTelemetry GenAI semantic conventions specification.
        """
        self._status: str | None = None
        self._validation_details: AttributeValidationDetails = {}
        self._spec = OtelGenAISpanSpecification()
        self._load_span(raw_span)
        self._validate()

    @cached_property
    def could_improve(self) -> bool:
        """Check if the span has recommended attributes that are missing.

        Returns:
            bool: True if there are any warnings indicating missing recommended attributes.
        """
        return self._has_validation_detail_of_severity(WARNING)

    @cached_property
    def has_deprecations(self) -> bool:
        """Check if the span has any deprecations.

        Returns:
            bool: True if there are any deprecated attributes in use.
        """
        return self._has_validation_detail_of_severity(INFO)

    @cached_property
    def has_errors(self) -> bool:
        """Check if the span has any errors.

        Returns:
            bool: True if the span has validation errors or failed to load.
        """
        return self._span is None or self._has_validation_detail_of_severity(ERROR)

    def extract_messages(self) -> list[dict[str, Any]]:
        """Extract message history from the span.

        Returns:
            A list of message dictionaries. Currently returns an empty list.
        """
        return []

    def extract_metadata(self) -> AttributeMetadata:
        """Extract common metadata fields from the span.

        Returns:
            AttributeMetadata: A structured object containing environment, model information,
                session/user IDs, timing data, and trace context extracted from the span.
        """
        return AttributeMetadata(
            environment=self.get_attribute_value("environment", "dev"),
            model_provider=self.get_attribute_value("gen_ai.provider.name"),
            model=self.get_attribute_value("gen_ai.request.model"),
            project_id=self.get_attribute_value("project.id"),
            session_id=self.extract_session_id(),
            start_time=self.get_span_root_value("start_time"),
            end_time=self.get_span_root_value("end_time"),
            span_id=format_hex_id(self.get_span_context_value("span_id")),
            span_parent_id=format_hex_id(self.get_span_root_value("parent_id")),
            trace_id=format_hex_id(self.get_span_context_value("trace_id")),
            user_id=self.get_attribute_value("user.id"),
        )

    def extract_session_id(self) -> str:
        """Extract session ID from the span, checking multiple possible attribute names.

        Returns:
            The session ID from either 'gen_ai.conversation.id' or 'session.id' attributes,
            or None if neither is present.
        """
        return self.get_attribute_value(
            "gen_ai.conversation.id"
        ) or self.get_attribute_value("session.id")

    def get_attribute_value(self, name: str, default_value: Any = None) -> Any:
        """Get a span attribute value by name.

        Args:
            name: The attribute name to retrieve.
            default_value: Value to return if attribute is not found. Defaults to None.

        Returns:
            The attribute value if found, otherwise the default value.
        """
        return (
            self._span.get("attributes", {}).get(name, default_value)
            if self._span
            else default_value
        )

    def get_span_context_value(self, name: str, default_value: Any = None) -> Any:
        """Get a span context value by name.

        Args:
            name: The context field name to retrieve (e.g., 'span_id', 'trace_id').
            default_value: Value to return if field is not found. Defaults to None.

        Returns:
            The context value if found, otherwise the default value.
        """
        return (
            self._span.get("context", {}).get(name, default_value)
            if self._span
            else default_value
        )

    def get_span_root_value(self, name: str, default_value: Any = None) -> Any:
        """Get a root-level span field value by name.

        Args:
            name: The root field name to retrieve (e.g., 'start_time', 'end_time', 'parent_id').
            default_value: Value to return if field is not found. Defaults to None.

        Returns:
            The root field value if found, otherwise the default value.
        """
        return self._span.get(name, default_value) if self._span else default_value

    def is_attribute_deprecated(self, name: str) -> bool:
        """Check if an attribute is deprecated.

        Args:
            name: The attribute name to check.

        Returns:
            bool: True if the attribute is deprecated.
        """
        return self._get_validation_detail_of_severity(name, INFO) is not None

    def is_attribute_invalid(self, name: str) -> bool:
        """Check if an attribute is invalid.

        Args:
            name: The attribute name to check.

        Returns:
            bool: True if the attribute has validation errors.
        """
        return self._get_validation_detail_of_severity(name, ERROR) is not None

    def is_attribute_missing(self, name: str) -> bool:
        """Check if an attribute is missing.

        Args:
            name: The attribute name to check.

        Returns:
            bool: True if the attribute is recommended but missing.
        """
        return self._get_validation_detail_of_severity(name, WARNING) is not None

    def is_attribute_valid(self, name: str) -> bool:
        """Check the validity of an attribute by name.

        Args:
            name: The attribute name to check.

        Returns:
            bool: True if the attribute has no validation errors.
        """
        return self._get_validation_detail_of_severity(name, ERROR) is None

    def _discern_severity(
        self, attribute_spec: AttributeSpec, key: str, value: Any
    ) -> AttributeSeverity | None:
        requirement = attribute_spec.get("requirement", OPTIONAL)

        # Handle missing values
        if value is None:
            if requirement == REQUIRED:
                logging.error(f"Attribute {key} is required but not set")
                return ERROR
            elif requirement == RECOMMENDED:
                return WARNING
            return None

        # Handle deprecated attributes
        if attribute_spec.get("deprecated", False):
            logging.warning(f"Attribute {key} is deprecated")
            return INFO

        # Validate type
        expected_type = attribute_spec.get("type")
        if expected_type and not is_valid_type(value, expected_type):
            logging.error(
                f"Attribute {key} has invalid type. Expected {expected_type}, got {type(value).__name__}"
            )
            return ERROR

        # Validate enum values
        valid_values = attribute_spec.get("valid_values", [])
        if valid_values:
            if isinstance(value, list):
                if not all(item in valid_values for item in value):
                    logging.error(
                        f"Attribute {key} has invalid values. Expected one of {valid_values}"
                    )
                    return ERROR
            elif value not in valid_values:
                logging.error(
                    f"Attribute {key} has invalid value '{value}'. Expected one of {valid_values}"
                )
                return ERROR

        return None

    def _get_validation_detail_of_severity(
        self, key: str, severity: AttributeSeverity
    ) -> AttributeValidationDetail | None:
        match = next(
            (
                v
                for v in self._validation_details.values()
                if v["name"] == key and v["severity"] == severity
            ),
            None,
        )

        return match

    def _has_validation_detail_of_severity(self, severity: AttributeSeverity) -> bool:
        return any(v["severity"] == severity for v in self._validation_details.values())

    def _is_operation_specific_attribute(
        self, attribute_name: str, operation: str
    ) -> bool:
        if attribute_name.startswith("gen_ai.evaluation."):
            return operation not in ["", None]

        if attribute_name.startswith("gen_ai.embeddings."):
            return operation != "embeddings"

        if attribute_name.startswith("gen_ai.tool."):
            return operation != "execute_tool"

        return False

    def _is_provider_specific_attribute(
        self, attribute_name: str, provider: str
    ) -> bool:
        provider_prefixes = {
            "aws.bedrock.": "aws.bedrock",
            "openai.": "openai",
        }

        for prefix, required_provider in provider_prefixes.items():
            if attribute_name.startswith(prefix):
                return provider != required_provider

        return False

    def _load_span(self, raw_span: Any) -> None:
        try:
            self._span = json.loads(raw_span) if isinstance(raw_span, str) else raw_span

            if (
                self.get_span_context_value("trace_id") is None
                or self.get_span_context_value("span_id") is None
            ):
                logging.error("Trace ID or span ID is not set")
                self._status = "Trace ID or span ID is not set"
                self._span = None

        except (json.JSONDecodeError, TypeError, ValueError) as e:
            logging.error(f"Failed to load span: {e}")
            self._status = "Invalid span JSON"
            self._span = None

    def _push_validation_detail(
        self, attribute_spec: AttributeSpec, key: str, value: Any
    ) -> None:
        requirement = attribute_spec.get("requirement", OPTIONAL)
        severity = self._discern_severity(attribute_spec, key, value)

        if severity:
            self._validation_details[key] = AttributeValidationDetail(
                name=attribute_spec["name"],
                requirement=requirement,
                severity=severity,
                type=attribute_spec.get("type"),
                deprecated=attribute_spec.get("deprecated"),
                deprecated_message=attribute_spec.get("deprecated_message"),
                valid_values=attribute_spec.get("valid_values"),
            )

    def _should_validate_attribute(self, attribute_name: str, attrs: dict) -> bool:
        provider = attrs.get("gen_ai.provider.name", "")
        operation = attrs.get("gen_ai.operation.name", "")

        return not self._is_provider_specific_attribute(
            attribute_name,
            provider,
        ) and not self._is_operation_specific_attribute(
            attribute_name,
            operation,
        )

    def _validate(self) -> None:
        if self._span is None:
            return

        try:
            attrs = self._span.get("attributes", {})
            for attribute_spec in self._spec:
                name = attribute_spec["name"]

                if not self._should_validate_attribute(name, attrs):
                    continue

                matching = [
                    key for key in attrs if key.startswith(f"{name}.") or key == name
                ]

                if not matching:
                    self._push_validation_detail(attribute_spec, name, None)
                else:
                    for key in matching:
                        self._push_validation_detail(attribute_spec, key, attrs[key])

        except (AttributeError, KeyError) as e:
            logging.error(f"Validation error: {e}")
            self._status = "An error occurred while validating the span"
