import json
import logging
from collections.abc import Iterator
from functools import cached_property
from typing import Any

from ..utils import format_hex_id, safe_json_loads
from .otlp_converter import convert_otlp_to_json
from .spec import OtelGenAISpecification
from .spec_message_input import OtelGenAIInputMessageSpecification
from .spec_message_input_part import (
    OtelGenAIInputMessageGenericPartSpecification,
    OtelGenAIInputMessageTextPartSpecification,
    OtelGenAIInputMessageToolCallRequestPartSpecification,
    OtelGenAIInputMessageToolCallResponsePartSpecification,
)
from .spec_message_output import OtelGenAIOutputMessageSpecification
from .spec_message_output_part import (
    OtelGenAIOutputMessageBlobPartSpecification,
    OtelGenAIOutputMessageFilePartSpecification,
    OtelGenAIOutputMessageGenericPartSpecification,
    OtelGenAIOutputMessageReasoningPartSpecification,
    OtelGenAIOutputMessageTextPartSpecification,
    OtelGenAIOutputMessageToolCallRequestPartSpecification,
    OtelGenAIOutputMessageToolCallResponsePartSpecification,
    OtelGenAIOutputMessageUriPartSpecification,
)
from .spec_span 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) in SDK JSON format.
                For OTLP wire format (from collectors/S3), use from_otlp() instead.
        """
        self._spec = OtelGenAISpanSpecification()
        self._span: dict[str, Any] | None = None
        self._status: str | None = None
        self._validation_details: AttributeValidationDetails = {}
        self._load_span(raw_span)
        self._validate()

    @classmethod
    def from_otlp(
        cls,
        wire_span: dict[str, Any],
        resource: dict[str, Any] | None = None,
        scope: dict[str, Any] | None = None,
    ) -> "OtelGenAISpecCheck":
        """Create a spec checker from an OTLP wire format span.

        Wire format spec: https://github.com/open-telemetry/opentelemetry-proto/blob/main/opentelemetry/proto/trace/v1/trace.proto

        Args:
            wire_span: A span dict from resourceSpans[].scopeSpans[].spans[].
            resource: Optional resource dict from resourceSpans[].resource.
            scope: Optional scope dict from resourceSpans[].scopeSpans[].scope.

        Returns:
            An OtelGenAISpecCheck instance with the converted span.
        """
        converted = convert_otlp_to_json(wire_span, resource, scope)
        return cls(converted)

    @classmethod
    def from_otlp_json(
        cls, json_span: dict[str, Any]
    ) -> Iterator["OtelGenAISpecCheck"]:
        """Create a spec checker for each wire span in an OTLP JSON export.

        Args:
            json_span: A span dict in OTLP JSON export format.

        Returns:
            An Iterator of OtelGenAISpecCheck instances with the converted spans.
        """
        for rs in json_span.get("resourceSpans", []):
            resource = rs.get("resource", {})
            for ss in rs.get("scopeSpans", []):
                scope = ss.get("scope", {})
                for span in ss.get("spans", []):
                    yield cls.from_otlp(span, resource, scope)

    @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 combining input and output messages.
        """
        try:
            messages = []
            for json_messages in [
                self.get_attribute_value("gen_ai.input.messages"),
                self.get_attribute_value("gen_ai.output.messages"),
            ]:
                if json_messages:
                    parsed_messages = safe_json_loads(json_messages, default=[])
                    if isinstance(parsed_messages, list):
                        messages.extend(parsed_messages)

            return messages
        except (AttributeError, KeyError, json.JSONDecodeError) as e:
            logging.error(f"Failed to extract messages: {e}")
            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_span(self) -> dict[str, Any] | None:
        """Return the loaded span."""
        return self._span

    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)

        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

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

        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

        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_part_spec_for_type(
        self, part_type: str, is_input: bool
    ) -> type[OtelGenAISpecification]:
        if is_input:
            type_map: dict[str, type[OtelGenAISpecification]] = {
                "text": OtelGenAIInputMessageTextPartSpecification,
                "tool_call": OtelGenAIInputMessageToolCallRequestPartSpecification,
                "tool_call_response": OtelGenAIInputMessageToolCallResponsePartSpecification,
            }
        else:
            type_map = {
                "text": OtelGenAIOutputMessageTextPartSpecification,
                "tool_call": OtelGenAIOutputMessageToolCallRequestPartSpecification,
                "tool_call_response": OtelGenAIOutputMessageToolCallResponsePartSpecification,
                "blob": OtelGenAIOutputMessageBlobPartSpecification,
                "file": OtelGenAIOutputMessageFilePartSpecification,
                "uri": OtelGenAIOutputMessageUriPartSpecification,
                "reasoning": OtelGenAIOutputMessageReasoningPartSpecification,
            }

        return type_map.get(
            part_type,
            OtelGenAIInputMessageGenericPartSpecification
            if is_input
            else OtelGenAIOutputMessageGenericPartSpecification,
        )

    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:
        self._span = safe_json_loads(raw_span, default=None)

        if self._span is None:
            logging.error("Failed to load span")
            self._status = "Invalid span JSON"
            return

        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

    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])

            self._validate_messages()

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

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

        attrs = self._span.get("attributes", {})

        for attribute_name, is_input in [
            ("gen_ai.input.messages", True),
            ("gen_ai.output.messages", False),
        ]:
            json_messages = attrs.get(attribute_name)
            if not json_messages:
                continue

            try:
                parsed_messages = safe_json_loads(json_messages, default=[])
                if not isinstance(parsed_messages, list):
                    logging.error(
                        f"Attribute {attribute_name} is not a list, got {type(parsed_messages).__name__}"
                    )
                    continue

                message_spec = (
                    OtelGenAIInputMessageSpecification()
                    if is_input
                    else OtelGenAIOutputMessageSpecification()
                )

                for msg_idx, message in enumerate(parsed_messages):
                    if not isinstance(message, dict):
                        logging.error(
                            f"Message at index {msg_idx} in {attribute_name} is not a dict"
                        )
                        continue

                    message_path = f"{attribute_name}.{msg_idx}"
                    self._validate_nested_object(message, message_spec, message_path)

                    parts = message.get("parts")
                    if parts is None:
                        continue

                    if not isinstance(parts, list):
                        logging.error(
                            f"Parts at {message_path} is not a list, got {type(parts).__name__}"
                        )
                        continue

                    for part_idx, part in enumerate(parts):
                        if not isinstance(part, dict):
                            logging.error(
                                f"Part at index {part_idx} in {message_path}.parts is not a dict"
                            )
                            continue

                        part_type = part.get("type")
                        if not isinstance(part_type, str):
                            part_spec_class = (
                                OtelGenAIInputMessageGenericPartSpecification
                                if is_input
                                else OtelGenAIOutputMessageGenericPartSpecification
                            )
                        else:
                            part_spec_class = self._get_part_spec_for_type(
                                part_type, is_input
                            )

                        part_spec = part_spec_class()
                        part_path = f"{message_path}.parts.{part_idx}"
                        self._validate_nested_object(part, part_spec, part_path)

            except json.JSONDecodeError as e:
                logging.error(f"Failed to parse {attribute_name}: {e}")
            except (AttributeError, KeyError) as e:
                logging.error(f"Error validating {attribute_name}: {e}")

    def _validate_nested_object(
        self, obj: dict[str, Any], spec: OtelGenAISpecification, base_path: str
    ) -> None:
        for attribute_spec in spec:
            name = attribute_spec["name"]
            path = f"{base_path}.{name}" if base_path else name
            value = obj.get(name)

            modified_spec = AttributeSpec(
                name=path,
                requirement=attribute_spec["requirement"],
                type=attribute_spec.get("type"),
                deprecated=attribute_spec.get("deprecated"),
                deprecated_message=attribute_spec.get("deprecated_message"),
                valid_values=attribute_spec.get("valid_values"),
            )
            self._push_validation_detail(modified_spec, path, value)
