"""Jaeger backend for distributed tracing logs."""

import asyncio
from dataclasses import dataclass
from typing import Any

from ..core.backend import BackendConfig, LoggingBackend, LogRecord
from ..formatters import JSONFormatter, LogFormatter

try:
    from jaeger_client import Config as JaegerConfig
    from jaeger_client.span import Span
    from jaeger_client.tracer import Tracer
    JAEGER_AVAILABLE = True
except ImportError:
    JAEGER_AVAILABLE = False
    JaegerConfig = None
    Tracer = None
    Span = None


@dataclass
class JaegerBackendConfig(BackendConfig):
    """Configuration for Jaeger backend."""

    enabled: bool = True
    timeout: float = 5.0
    retry_attempts: int = 3
    retry_delay: float = 1.0
    collector_endpoint: str = "http://localhost:14268/api/traces"
    service_name: str = "unified-logger"
    batch_size: int = 100


class JaegerBackend(LoggingBackend):
    """Backend that sends logs as traces to Jaeger.
    
    Note: This is a placeholder implementation. In a production environment,
    you would install and use jaeger-client: pip install jaeger-client
    """

    def __init__(
        self,
        name: str = "jaeger",
        config: JaegerBackendConfig | None = None,
        formatter: LogFormatter | None = None,
        **kwargs: Any
    ) -> None:
        """Initialize Jaeger backend.
        
        Args:
            name: Backend name
            config: Backend configuration
            formatter: Log formatter to use (defaults to JSON)
            **kwargs: Additional configuration
        """
        self.config = config or JaegerBackendConfig()
        super().__init__(name, self.config, **kwargs)

        self.formatter = formatter or JSONFormatter()
        self._tracer = None

        # For backward compatibility with tests
        self.service_name = kwargs.get('service_name', self.config.service_name)
        self.collector_endpoint = kwargs.get('collector_endpoint', self.config.collector_endpoint)
        self.timeout = kwargs.get('timeout', getattr(self.config, 'timeout', 30))

    async def connect(self) -> bool:
        """Connect to Jaeger collector."""
        if not JAEGER_AVAILABLE:
            return False

        try:
            # Initialize Jaeger tracer configuration
            config = JaegerConfig(
                config={
                    'sampler': {
                        'type': 'const',
                        'param': 1,
                    },
                    'logging': True,
                    'reporter_batch_size': 1,
                    'local_agent': {
                        'reporting_host': self.collector_endpoint.split('://')[1].split(':')[0],
                        'reporting_port': int(self.collector_endpoint.split(':')[-1].split('/')[0])
                            if ':' in self.collector_endpoint.split('://')[1]
                            else 14268,
                    },
                },
                service_name=self.service_name,
                validate=True,
            )

            # Initialize the tracer
            self._tracer = config.initialize_tracer()

            if self._tracer:
                self._connected = True
                return True
            else:
                return False

        except Exception as e:
            print(f"Failed to connect to Jaeger: {e}")
            return False

    async def disconnect(self) -> None:
        """Disconnect from Jaeger."""
        try:
            if self._tracer:
                # Close the tracer to flush any pending spans
                self._tracer.close()
                self._tracer = None

            self._connected = False
            return True

        except Exception as e:
            print(f"Failed to disconnect from Jaeger: {e}")
            return False

    async def send_log(self, record: LogRecord) -> bool:
        """Send log record as a trace span to Jaeger.
        
        Args:
            record: The log record to send
            
        Returns:
            True if successfully sent, False otherwise
        """
        if not self._tracer or not self._connected:
            return False

        try:
            # Convert log record to trace span
            import logging
            level_name = logging.getLevelName(record.level)
            with self._tracer.start_span(
                operation_name=f"{record.logger_name}.{level_name}",
                start_time=record.timestamp.timestamp()
            ) as span:
                import logging
                level_name = logging.getLevelName(record.level)
                span.set_tag('level', level_name)
                span.set_tag('logger', record.logger_name)
                span.set_tag('message', record.message)

                # Add extra fields as tags
                for key, value in record.extra.items():
                    span.set_tag(f'extra.{key}', str(value))

                # Add exception info if present
                if record.exc_info:
                    span.set_tag('error', True)
                    span.set_tag('error.kind', record.exc_info[0].__name__ if record.exc_info[0] else 'Unknown')
                    span.set_tag('error.message', str(record.exc_info[1]) if record.exc_info[1] else '')

                    # Format exception traceback
                    import traceback
                    exception_text = ''.join(traceback.format_exception(*record.exc_info))
                    span.log_kv({'event': 'error', 'error.object': exception_text})

                span.finish()

            return True

        except Exception:
            return False

    async def send_logs_batch(self, records: list[LogRecord]) -> list[bool]:
        """Send multiple log records as trace spans to Jaeger.
        
        Args:
            records: List of log records to send
            
        Returns:
            List of success/failure status for each record
        """
        # Send records individually in parallel
        tasks = [self.send_log(record) for record in records]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Convert exceptions to False
        return [
            result if isinstance(result, bool) else False
            for result in results
        ]

    async def health_check(self) -> bool:
        """Check Jaeger collector connectivity.
        
        Returns:
            True if collector is reachable, False otherwise
        """
        if not self._tracer or not self._connected:
            return False

        try:
            # Create a test span to verify tracer is working
            with self._tracer.start_span(
                operation_name="health_check",
                start_time=asyncio.get_event_loop().time()
            ) as span:
                span.set_tag('health_check', True)
                span.finish()

            return True

        except Exception as e:
            print(f"Health check failed: {e}")
            return False
