import asyncio
import aiohttp
import json
import logging
from threading import Lock
from typing import List, Dict, Any
from .types import TraceSpan
from .config import TracingConfig

logger = logging.getLogger(__name__)


class SafeJSONEncoder(json.JSONEncoder):
    """JSON encoder that handles non-serializable objects gracefully"""

    def default(self, obj: Any) -> Any:
        # Try to get a dict representation
        if hasattr(obj, 'dict'):
            return obj.dict()
        if hasattr(obj, 'to_dict'):
            return obj.to_dict()
        if hasattr(obj, '__dict__'):
            return {k: v for k, v in obj.__dict__.items() if not k.startswith('_')}
        # For other objects, return string representation
        return str(obj)


class TraceSender:
    """Sender for trace spans - groups spans by trace_id"""

    def __init__(self):
        self._lock = Lock()
        # Store spans grouped by trace_id
        self._traces: Dict[str, List[TraceSpan]] = {}

    def add_span(self, span: TraceSpan):
        """Add a span to its trace group"""
        with self._lock:
            trace_id = span.trace_id
            if trace_id not in self._traces:
                self._traces[trace_id] = []
            self._traces[trace_id].append(span)

    def flush_trace(self, trace_id: str):
        """Send all spans for a specific trace"""
        with self._lock:
            if trace_id not in self._traces:
                logger.warning(f"[TRACING] No spans found for trace {trace_id}")
                return
            spans = self._traces.pop(trace_id)

        if not spans:
            logger.warning(f"[TRACING] Empty spans list for trace {trace_id}")
            return

        logger.info(f"[TRACING] Flushing trace {trace_id} with {len(spans)} spans")
        self._send_trace_sync(trace_id, spans)

    def _send_trace_sync(self, trace_id: str, spans: List[TraceSpan]):
        """Synchronously send a complete trace"""
        url = TracingConfig.get_url()
        if not url:
            return

        # Build hierarchical trace structure
        trace_data = self._build_trace_structure(trace_id, spans)

        payload = {
            "project": TracingConfig.get_project(),
            "trace": trace_data
        }

        headers = {"Content-Type": "application/json"}
        api_key = TracingConfig.get_api_key()
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"

        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

        try:
            loop.run_until_complete(self._async_send(url, payload, headers))
        finally:
            loop.close()

    def _build_trace_structure(self, trace_id: str, spans: List[TraceSpan]) -> dict:
        """Build hierarchical trace structure from flat spans list"""
        # Create lookup by span_id
        span_map = {span.span_id: span.to_dict() for span in spans}

        # Find root spans (no parent or parent not in this trace)
        root_spans = []
        child_map: Dict[str, List[dict]] = {}

        for span in spans:
            span_dict = span_map[span.span_id]
            parent_id = span.parent_span_id

            if parent_id is None or parent_id not in span_map:
                root_spans.append(span_dict)
            else:
                if parent_id not in child_map:
                    child_map[parent_id] = []
                child_map[parent_id].append(span_dict)

        # Build tree recursively
        def attach_children(span_dict: dict):
            span_id = span_dict.get("span_id")
            if span_id in child_map:
                span_dict["children"] = child_map[span_id]
                for child in span_dict["children"]:
                    attach_children(child)

        for root in root_spans:
            attach_children(root)

        # Calculate trace-level metadata
        all_times = [s.start_time for s in spans if s.start_time]
        end_times = [s.end_time for s in spans if s.end_time]

        return {
            "trace_id": trace_id,
            "start_time": min(all_times) if all_times else None,
            "end_time": max(end_times) if end_times else None,
            "spans": root_spans,
            "span_count": len(spans)
        }

    async def _async_send(self, url: str, payload: dict, headers: dict):
        """Async HTTP POST"""
        try:
            logger.debug(f"[TRACING] Sending trace to {url}, spans: {payload.get('trace', {}).get('span_count', 0)}")
            # Use custom encoder to handle non-serializable objects
            json_data = json.dumps(payload, cls=SafeJSONEncoder)
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    url,
                    data=json_data,
                    headers=headers,
                    timeout=aiohttp.ClientTimeout(total=10)
                ) as resp:
                    if resp.status >= 400:
                        body = await resp.text()
                        logger.error(f"[TRACING] Failed to send trace: HTTP {resp.status}, {body}")
                    else:
                        logger.info(f"[TRACING] Trace sent successfully: HTTP {resp.status}")
        except Exception as e:
            logger.error(f"[TRACING] Exception sending trace: {e}")

    def flush(self):
        """Send all pending traces"""
        with self._lock:
            trace_ids = list(self._traces.keys())

        for trace_id in trace_ids:
            self.flush_trace(trace_id)

    def stop(self):
        """Stop the sender and flush remaining traces"""
        self.flush()