"""
Streaming support implementation for InferenceKit.
Handles streaming responses from providers with comprehensive features.
"""

import asyncio
import aiohttp
from typing import Dict, Any, List, Optional, Union, Callable, AsyncGenerator, Type
from datetime import datetime, timedelta
from uuid import UUID, uuid4
from enum import Enum
from collections import deque
from contextlib import asynccontextmanager
import json
import io
import struct
import time
import threading
import queue

from ..types import (
    ModelType,
    ProviderType,
    Request,
    Response,
    ModelConfig,
    BatchRequest,
    BatchResponse,
)
from ..exceptions import ProviderError, StreamingError, OperationTimeoutError
from ..core.router import Router


class StreamFormat(Enum):
    """Streaming data format enumeration"""
    TEXT = "text"
    JSON = "json"
    BINARY = "binary"
    MIXED = "mixed"


class StreamState(Enum):
    """Stream state enumeration"""
    IDLE = "idle"
    ACTIVE = "active"
    PAUSED = "paused"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class StreamChunk:
    """Represents a single chunk of streamed data"""

    def __init__(
        self,
        data: Union[str, Dict[str, Any], bytes],
        format: StreamFormat,
        metadata: Optional[Dict[str, Any]] = None,
        timestamp: Optional[datetime] = None,
        chunk_id: Optional[UUID] = None,
        total_chunks: Optional[int] = None
    ):
        self.data = data
        self.format = format
        self.metadata = metadata or {}
        self.timestamp = timestamp or datetime.utcnow()
        self.chunk_id = chunk_id or uuid4()
        self.total_chunks = total_chunks

    def __repr__(self) -> str:
        return f"StreamChunk(format={self.format}, size={len(str(self.data))})"


class StreamMetrics:
    """Streaming metrics and monitoring data"""

    def __init__(self):
        self.start_time: Optional[datetime] = None
        self.end_time: Optional[datetime] = None
        self.total_chunks: int = 0
        self.total_bytes: int = 0
        self.chunk_timestamps: List[datetime] = []
        self.bandwidth_samples: List[float] = []
        self.latency_samples: List[float] = []
        self.errors: List[Dict[str, Any]] = []

    def start(self) -> None:
        """Start metrics collection"""
        self.start_time = datetime.utcnow()
        self.end_time = None
        self.total_chunks = 0
        self.total_bytes = 0
        self.chunk_timestamps.clear()
        self.bandwidth_samples.clear()
        self.latency_samples.clear()
        self.errors.clear()

    def record_chunk(self, chunk: StreamChunk, bytes_transferred: int) -> None:
        """Record metrics for a chunk"""
        self.total_chunks += 1
        self.total_bytes += bytes_transferred
        self.chunk_timestamps.append(datetime.utcnow())

        # Calculate bandwidth (bytes per second)
        if len(self.chunk_timestamps) > 1:
            time_diff = (self.chunk_timestamps[-1] - self.chunk_timestamps[-2]).total_seconds()
            if time_diff > 0:
                bandwidth = bytes_transferred / time_diff
                self.bandwidth_samples.append(bandwidth)

        # Calculate latency
        if self.start_time:
            latency = (datetime.utcnow() - self.start_time).total_seconds()
            self.latency_samples.append(latency)

    def record_error(self, error: Exception, context: Dict[str, Any]) -> None:
        """Record an error"""
        self.errors.append({
            "timestamp": datetime.utcnow(),
            "error": str(error),
            "context": context
        })

    def stop(self) -> None:
        """Stop metrics collection"""
        self.end_time = datetime.utcnow()

    @property
    def duration(self) -> Optional[float]:
        """Get total duration in seconds"""
        if self.start_time and self.end_time:
            return (self.end_time - self.start_time).total_seconds()
        return None

    @property
    def average_bandwidth(self) -> Optional[float]:
        """Get average bandwidth in bytes per second"""
        if self.bandwidth_samples:
            return sum(self.bandwidth_samples) / len(self.bandwidth_samples)
        return None

    @property
    def average_latency(self) -> Optional[float]:
        """Get average latency in seconds"""
        if self.latency_samples:
            return sum(self.latency_samples) / len(self.latency_samples)
        return None

    def get_summary(self) -> Dict[str, Any]:
        """Get metrics summary"""
        return {
            "total_chunks": self.total_chunks,
            "total_bytes": self.total_bytes,
            "duration_seconds": self.duration,
            "average_bandwidth_bps": self.average_bandwidth,
            "average_latency_seconds": self.average_latency,
            "error_count": len(self.errors),
            "chunk_rate": self.total_chunks / self.duration if self.duration else 0,
            "start_time": self.start_time,
            "end_time": self.end_time
        }


class StreamContext:
    """Streaming context management"""

    def __init__(
        self,
        stream_id: UUID,
        request: Request,
        config: Dict[str, Any],
        metrics: StreamMetrics
    ):
        self.stream_id = stream_id
        self.request = request
        self.config = config
        self.metrics = metrics
        self.state = StreamState.IDLE
        self.creation_time = datetime.utcnow()
        self.timeout: Optional[datetime] = None
        self.cancellation_requested = False
        self.pause_queue: asyncio.Queue = asyncio.Queue()

    async def check_timeout(self) -> None:
        """Check if stream has timed out"""
        if self.timeout and datetime.utcnow() > self.timeout:
            raise OperationTimeoutError(f"Stream {self.stream_id} timed out")

    async def check_cancellation(self) -> None:
        """Check if stream has been cancelled"""
        if self.cancellation_requested:
            raise asyncio.CancelledError(f"Stream {self.stream_id} cancelled")

    async def pause(self) -> None:
        """Pause the stream"""
        self.state = StreamState.PAUSED
        await self.pause_queue.put(True)

    async def resume(self) -> None:
        """Resume the stream"""
        self.state = StreamState.ACTIVE
        await self.pause_queue.get()

    def cancel(self) -> None:
        """Request cancellation of the stream"""
        self.cancellation_requested = True

    def set_timeout(self, timeout_seconds: float) -> None:
        """Set stream timeout"""
        self.timeout = datetime.utcnow() + timedelta(seconds=timeout_seconds)


class StreamingManager:
    """Main streaming manager for handling streaming responses."""

    def __init__(self, **config):
        self.config = config
        self.active_streams: Dict[UUID, StreamContext] = {}
        self.stream_handlers: Dict[str, Callable] = {}
        self.max_concurrent_streams = config.get("max_concurrent_streams", 10)
        self.default_timeout = config.get("default_timeout", 60.0)
        self.buffer_size = config.get("buffer_size", 4096)
        self.flow_control_enabled = config.get("flow_control", True)
        self.metrics_enabled = config.get("metrics", True)

        # Initialize router for provider-specific handling
        self.router = Router()

    def register_stream_handler(self, provider_type: ProviderType, handler: Callable) -> None:
        """Register a stream handler for a specific provider type."""
        self.stream_handlers[provider_type.value] = handler

    async def create_stream(
        self,
        request: Request,
        handler: Optional[Callable] = None,
        timeout: Optional[float] = None,
        **kwargs
    ) -> AsyncGenerator[Union[str, Dict[str, Any], bytes], None]:
        """Create a new streaming response."""
        if request.id in self.active_streams:
            raise StreamingError(f"Stream already exists for request {request.id}")

        if timeout is None:
            timeout = self.default_timeout

        # Get appropriate handler
        handler = handler or self.stream_handlers.get(request.provider.value)
        if not handler:
            raise StreamingError(f"No stream handler available for provider {request.provider}")

        # Create streaming context
        metrics = StreamMetrics()
        context = StreamContext(
            stream_id=request.id,
            request=request,
            config=kwargs,
            metrics=metrics
        )
        context.set_timeout(timeout)

        self.active_streams[request.id] = context

        try:
            # Start streaming
            async for chunk in self._run_stream(request, handler, context):
                yield chunk

        finally:
            # Clean up when stream is done
            if request.id in self.active_streams:
                del self.active_streams[request.id]

    async def _run_stream(
        self,
        request: Request,
        handler: Callable,
        context: StreamContext
    ) -> AsyncGenerator[Union[str, Dict[str, Any], bytes], None]:
        """Run a streaming operation."""
        try:
            context.state = StreamState.ACTIVE
            context.metrics.start()

            async for raw_chunk in handler(request, context):
                # Process and format chunk
                chunk = await self._process_chunk(raw_chunk, request, context)

                # Check for cancellation and timeout
                await context.check_cancellation()
                await context.check_timeout()

                # Flow control - wait if paused
                if context.state == StreamState.PAUSED:
                    await context.pause_queue.get()

                # Flow control - buffer management
                if self.flow_control_enabled:
                    await self._apply_flow_control(context)

                # Record metrics
                bytes_transferred = len(str(chunk.data).encode('utf-8')) if isinstance(chunk.data, str) else len(chunk.data)
                context.metrics.record_chunk(chunk, bytes_transferred)

                yield chunk

            context.state = StreamState.COMPLETED

        except asyncio.CancelledError:
            context.state = StreamState.CANCELLED
            context.metrics.record_error(StreamingError("Stream cancelled"), {"stream_id": context.stream_id})
            raise

        except Exception as e:
            context.state = StreamState.FAILED
            context.metrics.record_error(e, {"stream_id": context.stream_id})
            raise StreamingError(f"Streaming failed: {e}")

        finally:
            context.metrics.stop()

    async def _process_chunk(
        self,
        raw_chunk: Any,
        request: Request,
        context: StreamContext
    ) -> StreamChunk:
        """Process and format a raw chunk."""

        # Determine format
        if isinstance(raw_chunk, str):
            format = StreamFormat.TEXT
            data = raw_chunk
        elif isinstance(raw_chunk, dict):
            format = StreamFormat.JSON
            data = raw_chunk
        elif isinstance(raw_chunk, bytes):
            format = StreamFormat.BINARY
            data = raw_chunk
        else:
            # Try to convert to string
            try:
                data = str(raw_chunk)
                format = StreamFormat.TEXT
            except Exception:
                raise StreamingError(f"Unsupported chunk type: {type(raw_chunk)}")

        # Create StreamChunk with metadata
        chunk = StreamChunk(
            data=data,
            format=format,
            metadata={
                "stream_id": context.stream_id,
                "request_id": request.id,
                "chunk_number": context.metrics.total_chunks + 1,
                "format": format.value,
                "timestamp": datetime.utcnow().isoformat()
            },
            total_chunks=None  # Could be set by provider if known
        )

        return chunk

    async def _apply_flow_control(self, context: StreamContext) -> None:
        """Apply flow control to prevent overwhelming the consumer."""
        if context.metrics.total_chunks % 10 == 0:  # Throttle every 10 chunks
            await asyncio.sleep(0.01)  # Small delay to prevent overwhelming

    async def stream_generate(
        self,
        prompt: str,
        model: str,
        provider: ProviderType,
        **kwargs
    ) -> AsyncGenerator[Union[str, Dict[str, Any], bytes], None]:
        """Stream text generation."""
        request = Request(
            model=model,
            provider=provider,
            type=ModelType.TEXT_GENERATION,
            prompt=prompt,
            parameters=kwargs,
        )

        return self.create_stream(request, **kwargs)

    async def stream_chat(
        self,
        messages: List[Dict[str, Any]],
        model: str,
        provider: ProviderType,
        **kwargs
    ) -> AsyncGenerator[Union[str, Dict[str, Any], bytes], None]:
        """Stream chat conversation."""
        request = Request(
            model=model,
            provider=provider,
            type=ModelType.CHAT,
            prompt="\n".join([
                f"{msg.get('role', 'user')}: {msg.get('content', '')}"
                for msg in messages
            ]),
            parameters=kwargs,
        )

        return self.create_stream(request, **kwargs)

    async def stream_embeddings(
        self,
        text: str,
        model: str,
        provider: ProviderType,
        **kwargs
    ) -> AsyncGenerator[Union[str, Dict[str, Any], bytes], None]:
        """Stream embedding generation."""
        request = Request(
            model=model,
            provider=provider,
            type=ModelType.EMBEDDING,
            prompt=text,
            parameters=kwargs,
        )

        return self.create_stream(request, **kwargs)

    def get_active_streams(self) -> List[UUID]:
        """Get list of active stream IDs."""
        return list(self.active_streams.keys())

    async def cancel_stream(self, stream_id: UUID) -> bool:
        """Cancel an active stream."""
        if stream_id not in self.active_streams:
            return False

        context = self.active_streams[stream_id]
        context.cancel()

        if stream_id in self.active_streams:
            del self.active_streams[stream_id]
        return True

    def get_stream_stats(self) -> Dict[str, Any]:
        """Get streaming statistics."""
        return {
            "active_streams": len(self.active_streams),
            "max_concurrent_streams": self.max_concurrent_streams,
            "default_timeout": self.default_timeout,
            "buffer_size": self.buffer_size,
            "flow_control_enabled": self.flow_control_enabled,
            "metrics_enabled": self.metrics_enabled,
        }

    async def cleanup_streams(self) -> None:
        """Clean up completed or failed streams."""
        completed_tasks = []
        for stream_id, context in self.active_streams.items():
            if context.state in [StreamState.COMPLETED, StreamState.FAILED, StreamState.CANCELLED]:
                completed_tasks.append(stream_id)

        for stream_id in completed_tasks:
            del self.active_streams[stream_id]

    def get_stream_metrics(self, stream_id: UUID) -> Optional[Dict[str, Any]]:
        """Get metrics for a specific stream."""
        if stream_id in self.active_streams:
            context = self.active_streams[stream_id]
            return context.metrics.get_summary()
        return None

    def get_all_metrics(self) -> Dict[str, Any]:
        """Get metrics for all active streams."""
        metrics = {}
        for stream_id, context in self.active_streams.items():
            metrics[str(stream_id)] = context.metrics.get_summary()
        return metrics


# Provider-specific stream handlers
async def openai_stream_handler(request: Request, context: StreamContext) -> AsyncGenerator[Any, None]:
    """OpenAI stream handler."""
    from inferencekit.providers.openai import OpenAIProvider

    provider = OpenAIProvider({})
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{provider.base_url}/v1/chat/completions",
            headers=provider.get_headers(),
            json={
                "model": request.model,
                "messages": [{"role": "user", "content": request.prompt}],
                "stream": True,
                **request.parameters
            },
            timeout=aiohttp.ClientTimeout(total=context.config.get("timeout", 60))
        ) as response:
            async for chunk in response.content:
                if chunk:
                    # Parse streaming response
                    try:
                        data = json.loads(chunk.decode('utf-8'))
                        if 'choices' in data and 'delta' in data['choices'][0]:
                            yield data['choices'][0]['delta'].get('content', '')
                    except json.JSONDecodeError:
                        continue


async def anthropic_stream_handler(request: Request, context: StreamContext) -> AsyncGenerator[Any, None]:
    """Anthropic stream handler."""
    from inferencekit.providers.anthropic import AnthropicProvider

    provider = AnthropicProvider({})
    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{provider.base_url}/v1/messages",
            headers=provider.get_headers(),
            json={
                "model": request.model,
                "max_tokens": request.parameters.get("max_tokens", 1024),
                "messages": [{"role": "user", "content": request.prompt}],
                "stream": True,
                **request.parameters
            },
            timeout=aiohttp.ClientTimeout(total=context.config.get("timeout", 60))
        ) as response:
            async for chunk in response.content:
                if chunk:
                    try:
                        data = json.loads(chunk.decode('utf-8'))
                        if 'delta' in data:
                            yield data['delta'].get('content', '')
                    except json.JSONDecodeError:
                        continue


async def local_stream_handler(request: Request, context: StreamContext) -> AsyncGenerator[Any, None]:
    """Local model stream handler."""
    # For local models, we can simulate streaming by chunking the response
    from inferencekit.providers.local import LocalProvider

    provider = LocalProvider({})

    # Generate response (this would be async for real local models)
    response = await provider.generate(request.prompt, model=request.model, **request.parameters)

    # Simulate streaming by yielding chunks
    chunk_size = 100  # Characters per chunk
    text = response.content

    for i in range(0, len(text), chunk_size):
        await asyncio.sleep(0.05)  # Simulate streaming delay
        yield text[i:i + chunk_size]


# Initialize default handlers
streaming_manager = StreamingManager()
streaming_manager.register_stream_handler(ProviderType.OPENAI, openai_stream_handler)
streaming_manager.register_stream_handler(ProviderType.ANTHROPIC, anthropic_stream_handler)
streaming_manager.register_stream_handler(ProviderType.LOCAL, local_stream_handler)


# Helper functions for common streaming operations
async def stream_text_with_progress(
    prompt: str,
    model: str,
    provider: ProviderType,
    progress_callback: Optional[Callable[[float], None]] = None,
    **kwargs
) -> AsyncGenerator[str, None]:
    """Stream text with progress updates."""
    async for chunk in streaming_manager.stream_generate(prompt, model, provider, **kwargs):
        if progress_callback and hasattr(chunk, 'metadata'):
            progress = chunk.metadata.get('chunk_number', 0) / 100  # Simplified progress
            progress_callback(progress)

        if isinstance(chunk, str):
            yield chunk
        elif isinstance(chunk, StreamChunk):
            yield chunk.data


@asynccontextmanager
async def streaming_context(
    prompt: str,
    model: str,
    provider: ProviderType,
    **kwargs
):
    """Context manager for streaming operations."""
    stream = None
    try:
        stream = streaming_manager.stream_generate(prompt, model, provider, **kwargs)
        yield stream
    finally:
        if stream:
            # Ensure stream is properly closed
            try:
                async for _ in stream:
                    pass
            except Exception:
                pass


# Sync streaming wrapper for compatibility
class SyncStreamingWrapper:
    """Wrapper for sync streaming compatibility."""

    def __init__(self, async_generator: AsyncGenerator):
        self.async_generator = async_generator
        self._loop = asyncio.get_event_loop()

    def __aiter__(self):
        return self.async_generator.__aiter__()

    def __anext__(self):
        return self.async_generator.__anext__()

    def __iter__(self):
        """Make it iterable in sync contexts"""
        return self._sync_iter()

    def _sync_iter(self):
        """Sync iterator"""
        iterator = self.async_generator.__aiter__()
        while True:
            try:
                result = self._loop.run_until_complete(iterator.__anext__())
                yield result
            except StopAsyncIteration:
                break
            except Exception as e:
                raise StopIteration from e