import bisect
import datetime
from collections import defaultdict
from typing import Any, Dict, List, Optional, Tuple
import json
from moto.core import BaseBackend, BackendDict, BaseModel
from moto.core.exceptions import AWSError
from .exceptions import BadSegmentException


class TelemetryRecords(BaseModel):
    def __init__(
        self,
        instance_id: str,
        hostname: str,
        resource_arn: str,
        records: List[Dict[str, Any]],
    ):
        self.instance_id = instance_id
        self.hostname = hostname
        self.resource_arn = resource_arn
        self.records = records

    @classmethod
    def from_json(cls, src: Dict[str, Any]) -> "TelemetryRecords":  # type: ignore[misc]
        instance_id = src.get("EC2InstanceId", None)
        hostname = src.get("Hostname")
        resource_arn = src.get("ResourceARN")
        telemetry_records = src["TelemetryRecords"]

        return cls(instance_id, hostname, resource_arn, telemetry_records)  # type: ignore


# https://docs.aws.amazon.com/xray/latest/devguide/xray-api-segmentdocuments.html
class TraceSegment(BaseModel):
    def __init__(
        self,
        name: str,
        segment_id: str,
        trace_id: str,
        start_time: float,
        raw: Any,
        end_time: Optional[float] = None,
        in_progress: bool = False,
        service: Any = None,
        user: Any = None,
        origin: Any = None,
        parent_id: Any = None,
        http: Any = None,
        aws: Any = None,
        metadata: Any = None,
        annotations: Any = None,
        subsegments: Any = None,
        **kwargs: Any
    ):
        self.name = name
        self.id = segment_id
        self.trace_id = trace_id
        self._trace_version: Optional[int] = None
        self._original_request_start_time: Optional[datetime.datetime] = None
        self._trace_identifier = None
        self.start_time = start_time
        self._start_date: Optional[datetime.datetime] = None
        self.end_time = end_time
        self._end_date: Optional[datetime.datetime] = None
        self.in_progress = in_progress
        self.service = service
        self.user = user
        self.origin = origin
        self.parent_id = parent_id
        self.http = http
        self.aws = aws
        self.metadata = metadata
        self.annotations = annotations
        self.subsegments = subsegments
        self.misc = kwargs

        # Raw json string
        self.raw = raw

    def __lt__(self, other: Any) -> bool:
        return self.start_date < other.start_date

    @property
    def trace_version(self) -> int:
        if self._trace_version is None:
            self._trace_version = int(self.trace_id.split("-", 1)[0])
        return self._trace_version

    @property
    def request_start_date(self) -> datetime.datetime:
        if self._original_request_start_time is None:
            start_time = int(self.trace_id.split("-")[1], 16)
            self._original_request_start_time = datetime.datetime.fromtimestamp(
                start_time
            )
        return self._original_request_start_time

    @property
    def start_date(self) -> datetime.datetime:
        if self._start_date is None:
            self._start_date = datetime.datetime.fromtimestamp(self.start_time)
        return self._start_date

    @property
    def end_date(self) -> datetime.datetime:
        if self._end_date is None:
            self._end_date = datetime.datetime.fromtimestamp(self.end_time)  # type: ignore
        return self._end_date

    @classmethod
    def from_dict(cls, data: Dict[str, Any], raw: Any) -> "TraceSegment":  # type: ignore[misc]
        # Check manditory args
        if "id" not in data:
            raise BadSegmentException(code="MissingParam", message="Missing segment ID")
        seg_id = data["id"]
        data["segment_id"] = seg_id  # Just adding this key for future convenience

        for arg in ("name", "trace_id", "start_time"):
            if arg not in data:
                raise BadSegmentException(
                    seg_id=seg_id, code="MissingParam", message="Missing segment ID"
                )

        if "end_time" not in data and "in_progress" not in data:
            raise BadSegmentException(
                seg_id=seg_id,
                code="MissingParam",
                message="Missing end_time or in_progress",
            )
        if "end_time" not in data and data["in_progress"] == "false":
            raise BadSegmentException(
                seg_id=seg_id, code="MissingParam", message="Missing end_time"
            )

        return cls(raw=raw, **data)


class SegmentCollection(object):
    def __init__(self) -> None:
        self._traces: Dict[str, Dict[str, Any]] = defaultdict(self._new_trace_item)

    @staticmethod
    def _new_trace_item() -> Dict[str, Any]:  # type: ignore[misc]
        return {
            "start_date": datetime.datetime(1970, 1, 1),
            "end_date": datetime.datetime(1970, 1, 1),
            "finished": False,
            "trace_id": None,
            "segments": [],
        }

    def put_segment(self, segment: Any) -> None:
        # insert into a sorted list
        bisect.insort_left(self._traces[segment.trace_id]["segments"], segment)

        # Get the last segment (takes into account incorrect ordering)
        # and if its the last one, mark trace as complete
        if self._traces[segment.trace_id]["segments"][-1].end_time is not None:
            self._traces[segment.trace_id]["finished"] = True

            start_time = self._traces[segment.trace_id]["segments"][0].start_date
            end_time = self._traces[segment.trace_id]["segments"][-1].end_date
            self._traces[segment.trace_id]["start_date"] = start_time
            self._traces[segment.trace_id]["end_date"] = end_time
            self._traces[segment.trace_id]["trace_id"] = segment.trace_id
            # Todo consolidate trace segments into a trace.
            # not enough working knowledge of xray to do this

    def summary(
        self, start_time: str, end_time: str, filter_expression: Any = None
    ) -> Dict[str, Any]:
        # This beast https://docs.aws.amazon.com/xray/latest/api/API_GetTraceSummaries.html#API_GetTraceSummaries_ResponseSyntax
        if filter_expression is not None:
            raise AWSError(
                "Not implemented yet - moto",
                exception_type="InternalFailure",
                status=500,
            )

        summaries = []

        for tid, trace in self._traces.items():
            if (
                trace["finished"]
                and start_time < trace["start_date"]
                and trace["end_date"] < end_time
            ):
                duration = int(
                    (trace["end_date"] - trace["start_date"]).total_seconds()
                )
                # this stuff is mostly guesses, refer to TODO above
                has_error = any(["error" in seg.misc for seg in trace["segments"]])
                has_fault = any(["fault" in seg.misc for seg in trace["segments"]])
                has_throttle = any(
                    ["throttle" in seg.misc for seg in trace["segments"]]
                )

                # Apparently all of these options are optional
                summary_part = {
                    "Annotations": {},  # Not implemented yet
                    "Duration": duration,
                    "HasError": has_error,
                    "HasFault": has_fault,
                    "HasThrottle": has_throttle,
                    "Http": {},  # Not implemented yet
                    "Id": tid,
                    "IsParital": False,  # needs lots more work to work on partials
                    "ResponseTime": 1,  # definitely 1ms resposnetime
                    "ServiceIds": [],  # Not implemented yet
                    "Users": {},  # Not implemented yet
                }
                summaries.append(summary_part)

        result = {
            "ApproximateTime": int(
                (
                    datetime.datetime.now() - datetime.datetime(1970, 1, 1)
                ).total_seconds()
            ),
            "TracesProcessedCount": len(summaries),
            "TraceSummaries": summaries,
        }

        return result

    def get_trace_ids(
        self, trace_ids: List[str]
    ) -> Tuple[List[Dict[str, Any]], List[str]]:
        traces = []
        unprocessed = []

        # Its a default dict
        existing_trace_ids = list(self._traces.keys())
        for trace_id in trace_ids:
            if trace_id in existing_trace_ids:
                traces.append(self._traces[trace_id])
            else:
                unprocessed.append(trace_id)

        return traces, unprocessed


class XRayBackend(BaseBackend):
    def __init__(self, region_name: str, account_id: str):
        super().__init__(region_name, account_id)
        self._telemetry_records: List[TelemetryRecords] = []
        self._segment_collection = SegmentCollection()

    @staticmethod
    def default_vpc_endpoint_service(
        service_region: str, zones: List[str]
    ) -> List[Dict[str, str]]:
        """Default VPC endpoint service."""
        return BaseBackend.default_vpc_endpoint_service_factory(
            service_region, zones, "xray"
        )

    def add_telemetry_records(self, src: Any) -> None:
        self._telemetry_records.append(TelemetryRecords.from_json(src))

    def process_segment(self, doc: Any) -> None:
        try:
            data = json.loads(doc)
        except ValueError:
            raise BadSegmentException(code="JSONFormatError", message="Bad JSON data")

        try:
            # Get Segment Object
            segment = TraceSegment.from_dict(data, raw=doc)
        except ValueError:
            raise BadSegmentException(code="JSONFormatError", message="Bad JSON data")

        try:
            # Store Segment Object
            self._segment_collection.put_segment(segment)
        except Exception as err:
            raise BadSegmentException(
                seg_id=segment.id, code="InternalFailure", message=str(err)
            )

    def get_trace_summary(
        self, start_time: str, end_time: str, filter_expression: Any
    ) -> Dict[str, Any]:
        return self._segment_collection.summary(start_time, end_time, filter_expression)

    def get_trace_ids(self, trace_ids: List[str]) -> Dict[str, Any]:
        traces, unprocessed_ids = self._segment_collection.get_trace_ids(trace_ids)

        result: Dict[str, Any] = {"Traces": [], "UnprocessedTraceIds": unprocessed_ids}

        for trace in traces:
            segments = []
            for segment in trace["segments"]:
                segments.append({"Id": segment.id, "Document": segment.raw})

            result["Traces"].append(
                {
                    "Duration": int(
                        (trace["end_date"] - trace["start_date"]).total_seconds()
                    ),
                    "Id": trace["trace_id"],
                    "Segments": segments,
                }
            )

        return result


xray_backends = BackendDict(XRayBackend, "xray")
