"""
Streaming event schemas and coordinator for Responses API.

Handles Server-Sent Events (SSE) streaming with proper event types:
- response.reasoning_summary_text.delta
- response.output_text.delta
- response.function_call_arguments.delta
- response.done
"""

from enum import Enum
from typing import Literal, AsyncGenerator
import json
import time

from src.core.models import CustomBaseModel
from src.responses.schemas import Response
from src.agents.registry import ModelRegistry


# ============================================================================
# Stream Event Types
# ============================================================================

class StreamEventType(str, Enum):
    """Streaming event types for Responses API."""
    REASONING_SUMMARY_DELTA = "response.reasoning_summary_text.delta"
    REASONING_SUMMARY_DONE = "response.reasoning_summary_text.done"
    OUTPUT_TEXT_DELTA = "response.output_text.delta"
    OUTPUT_TEXT_DONE = "response.output_text.done"
    FUNCTION_CALL_DELTA = "response.function_call_arguments.delta"
    FUNCTION_CALL_DONE = "response.function_call_arguments.done"
    RESPONSE_DONE = "response.done"
    ERROR = "error"


# ============================================================================
# Stream Event Schemas
# ============================================================================

class ReasoningSummaryDelta(CustomBaseModel):
    """Reasoning summary text delta event."""
    event: Literal[StreamEventType.REASONING_SUMMARY_DELTA] = StreamEventType.REASONING_SUMMARY_DELTA
    delta: str


class ReasoningSummaryDone(CustomBaseModel):
    """Reasoning summary completion event."""
    event: Literal[StreamEventType.REASONING_SUMMARY_DONE] = StreamEventType.REASONING_SUMMARY_DONE


class OutputTextDelta(CustomBaseModel):
    """Output text delta event."""
    event: Literal[StreamEventType.OUTPUT_TEXT_DELTA] = StreamEventType.OUTPUT_TEXT_DELTA
    delta: str


class OutputTextDone(CustomBaseModel):
    """Output text completion event."""
    event: Literal[StreamEventType.OUTPUT_TEXT_DONE] = StreamEventType.OUTPUT_TEXT_DONE


class FunctionCallDelta(CustomBaseModel):
    """Function call arguments delta event."""
    event: Literal[StreamEventType.FUNCTION_CALL_DELTA] = StreamEventType.FUNCTION_CALL_DELTA
    delta: str
    name: str | None = None  # Only in first chunk


class FunctionCallDone(CustomBaseModel):
    """Function call completion event."""
    event: Literal[StreamEventType.FUNCTION_CALL_DONE] = StreamEventType.FUNCTION_CALL_DONE


class ResponseDone(CustomBaseModel):
    """Response completion event with full response."""
    event: Literal[StreamEventType.RESPONSE_DONE] = StreamEventType.RESPONSE_DONE
    response: Response


class ErrorEvent(CustomBaseModel):
    """Error event."""
    event: Literal[StreamEventType.ERROR] = StreamEventType.ERROR
    error: dict


# Union type for all stream events
StreamEvent = (
    ReasoningSummaryDelta |
    ReasoningSummaryDone |
    OutputTextDelta |
    OutputTextDone |
    FunctionCallDelta |
    FunctionCallDone |
    ResponseDone |
    ErrorEvent
)


# ============================================================================
# Streaming Coordinator
# ============================================================================

class StreamingCoordinator:
    """
    Coordinates streaming from agents to SSE format.

    Real production streaming logic that handles:
    1. Reasoning summary chunks
    2. Function call arguments
    3. Output text chunks
    4. Error handling
    5. Final response event
    """

    async def stream_response_with_steward(
        self,
        request: "ResponseRequest"  # type: ignore # Forward reference
    ) -> AsyncGenerator[StreamEvent, None]:
        """
        Stream response with Steward preprocessing (Phase 2 flow).

        Streams in order:
        1. Steward's analysis as reasoning summary
        2. Tatlock's response as output text

        Args:
            request: Response request

        Yields:
            StreamEvent: Stream of SSE events
        """
        from src.responses.service import _calculate_usage, generate_id, _conversation_history
        from src.core.preprocessing import preprocess_request
        from src.core.tool_tracking import ToolCallTracker
        from src.responses.schemas import MessageOutputItem, ReasoningOutputItem, OutputTextContent
        from src.agents.tatlock import TatlockAgent
        import asyncio

        output_items = []

        try:
            # Get or generate conversation ID
            conversation_id = await _conversation_history.get_conversation_id(request)

            # Extract user message and conversation history
            user_message = ""
            for msg in reversed(request.input):
                if msg.get("role") == "user":
                    user_message = msg.get("content", "")
                    break

            conversation_history = request.input[:-1] if len(request.input) > 1 else []

            # Phase 1: Steward preprocessing
            enriched = await preprocess_request(
                user_message,
                conversation_history=conversation_history,
                conversation_id=conversation_id,
            )

            # Stream Steward's analysis as reasoning summary
            steward_lines = enriched.steward_reasoning.split('\n')
            for line in steward_lines:
                if line.strip():
                    yield ReasoningSummaryDelta(delta=line + "\n")
                    await asyncio.sleep(0.05)

            yield ReasoningSummaryDone()

            # Add Steward reasoning to output items
            reasoning_item = ReasoningOutputItem(
                id=f"reasoning_{generate_id()}",
                summary=[
                    "🎩 Steward's Analysis:",
                    enriched.steward_reasoning,
                ],
                status="completed"
            )
            output_items.append(reasoning_item)

            # Phase 2: Initialize tool tracker
            tracker = ToolCallTracker(
                recommended_capabilities=enriched.recommendation.recommended_capabilities,
                conversation_id=conversation_id,
            )

            # Phase 3: Stream Tatlock's response with scoped tools
            tatlock = TatlockAgent()
            tatlock_response_parts = []

            async for chunk in tatlock.run_with_scoped_tools_stream(
                user_message=user_message,
                steward_note=enriched.steward_note,
                scoped_tools=enriched.scoped_tools,
                message_history=conversation_history,
                tool_tracker=tracker,
            ):
                tatlock_response_parts.append(chunk)
                yield OutputTextDelta(delta=chunk)

            yield OutputTextDone()

            # Combine response for output item
            tatlock_response = "".join(tatlock_response_parts)

            # Add Tatlock message to output items
            message_item = MessageOutputItem(
                id=f"msg_{generate_id()}",
                role="assistant",
                content=[OutputTextContent(
                    type="output_text",
                    text=tatlock_response,
                    annotations=[]
                )],
                status="completed"
            )
            output_items.append(message_item)

            # Phase 4: Finalize tool tracking
            await tracker.finalize()

            # Calculate usage and build final response
            usage = _calculate_usage(request.input, output_items)

            final_response = Response(
                id=f"resp_{generate_id()}",
                created_at=int(time.time()),
                model=request.model,
                status="completed",
                output=output_items,
                usage=usage
            )

            # Track conversation history
            await _conversation_history.add_response(conversation_id, final_response)

            yield ResponseDone(response=final_response)

        except Exception as e:
            # Stream error event
            yield self._create_error_event(e)

    async def stream_response(
        self,
        request: "ResponseRequest"  # type: ignore # Forward reference
    ) -> AsyncGenerator[StreamEvent, None]:
        """
        Coordinate streaming from agent to SSE events.

        Args:
            request: Response request

        Yields:
            StreamEvent: Stream of SSE events

        Example SSE output:
            event: response.reasoning_summary_text.delta
            data: {"delta": "Analyzing..."}

            event: response.output_text.delta
            data: {"delta": "Lorem"}

            event: response.done
            data: {"response": {...}}
        """
        from src.responses.service import _calculate_usage, generate_id
        import asyncio

        output_items = []
        last_message_text = ""  # Track last streamed message text to compute deltas

        try:
            # Strip pipeline prefix if present (e.g., "pipeline.model" -> "model")
            model_id = request.model
            if "." in model_id:
                model_id = model_id.split(".", 1)[1]

            # Get agent for model
            agent = ModelRegistry.get_agent(model_id)
            # Stream from agent
            async for item in agent.generate_response(
                messages=request.input,
                reasoning=request.reasoning,
                tools=request.tools,
                temperature=request.temperature,
                max_tokens=request.max_output_tokens,
                stop=request.stop,
            ):
                output_items.append(item)

                # Stream based on item type
                if item.type == "reasoning":
                    # Stream reasoning summary
                    for step in item.data.get("summary", []):
                        yield ReasoningSummaryDelta(delta=step + " ")
                        await asyncio.sleep(0.05)  # Simulate typing

                    yield ReasoningSummaryDone()

                elif item.type == "function_call":
                    # Stream function call arguments
                    # First chunk includes name
                    yield FunctionCallDelta(
                        name=item.data["name"],
                        delta=""
                    )

                    # Stream arguments in chunks
                    args = item.data["arguments"]
                    chunk_size = 20
                    for i in range(0, len(args), chunk_size):
                        yield FunctionCallDelta(
                            delta=args[i:i+chunk_size]
                        )
                        await asyncio.sleep(0.03)

                    yield FunctionCallDone()

                elif item.type == "message":
                    # Get current accumulated text from agent
                    current_text = item.data["content"][0]["text"]

                    # Only stream the NEW text (delta) since last update
                    if current_text.startswith(last_message_text):
                        # Extract only the new portion
                        delta_text = current_text[len(last_message_text):]

                        if delta_text:
                            # Stream the delta text in chunks while preserving formatting
                            # (newlines, markdown, code blocks, etc.)
                            chunk_size = 50  # characters per chunk

                            for i in range(0, len(delta_text), chunk_size):
                                chunk = delta_text[i:i+chunk_size]

                                # Check stop sequences on full accumulated text
                                stop_found, text_before_stop = self._check_stop_sequence(
                                    current_text,
                                    request.stop
                                )

                                if stop_found:
                                    # Only emit remaining delta before stop
                                    remaining = text_before_stop[len(last_message_text):]
                                    if remaining:
                                        yield OutputTextDelta(delta=remaining)
                                    yield OutputTextDone()
                                    break

                                # Check max tokens on full text
                                output_tokens = self._count_tokens_approx(current_text)
                                if self._check_max_tokens(output_tokens, request.max_output_tokens):
                                    yield OutputTextDone()
                                    break

                                # Normal streaming of delta chunk (preserves all formatting)
                                yield OutputTextDelta(delta=chunk)
                                await asyncio.sleep(0.02)  # Shorter delay since chunks are larger

                        # Update tracking variable
                        last_message_text = current_text

                    # If this is the final message (status=completed), ensure we send done
                    if item.data.get("status") == "completed":
                        yield OutputTextDone()

            # Final response.done event with complete response
            usage = _calculate_usage(request.input, output_items)

            final_response = Response(
                id=f"resp_{generate_id()}",
                created_at=int(time.time()),
                model=request.model,
                status="completed",
                output=self._convert_output_items(output_items),
                usage=usage
            )

            # Track conversation history (import here to avoid circular dependency)
            from src.responses.service import _conversation_history
            conversation_id = await _conversation_history.get_conversation_id(request)
            await _conversation_history.add_response(conversation_id, final_response)

            yield ResponseDone(response=final_response)

        except Exception as e:
            # Stream error event
            yield self._create_error_event(e)

    def _convert_output_items(self, items: list) -> list:
        """Convert agent OutputItem objects to schema OutputItem objects."""
        from src.responses.schemas import (
            MessageOutputItem,
            ReasoningOutputItem,
            FunctionCallOutputItem,
            OutputTextContent,
        )

        converted = []
        for item in items:
            if item.type == "message":
                converted.append(MessageOutputItem(
                    id=item.id,
                    content=[OutputTextContent(**c) for c in item.data["content"]],
                    status=item.data.get("status", "completed")
                ))
            elif item.type == "reasoning":
                converted.append(ReasoningOutputItem(
                    id=item.id,
                    summary=item.data["summary"],
                    status=item.data.get("status", "completed")
                ))
            elif item.type == "function_call":
                converted.append(FunctionCallOutputItem(
                    id=item.id,
                    name=item.data["name"],
                    arguments=item.data["arguments"],
                    status=item.data.get("status", "completed")
                ))

        return converted

    def _create_error_event(self, error: Exception) -> ErrorEvent:
        """Create error event from exception."""
        from src.core.exceptions import (
            RateLimitError,
            ContextLengthError,
            AppException,
        )

        if isinstance(error, RateLimitError):
            error_type = "rate_limit_exceeded"
            code = 429
        elif isinstance(error, ContextLengthError):
            error_type = "context_length_exceeded"
            code = 400
        elif isinstance(error, AppException):
            error_type = "app_error"
            code = error.status_code
        else:
            error_type = "internal_error"
            code = 500

        return ErrorEvent(
            error={
                "type": error_type,
                "message": str(error),
                "code": code
            }
        )

    def _check_stop_sequence(
        self,
        accumulated_text: str,
        stop_sequences: list[str] | None
    ) -> tuple[bool, str]:
        """
        Check if any stop sequence is encountered.

        Args:
            accumulated_text: Text accumulated so far
            stop_sequences: List of stop sequences to check

        Returns:
            tuple: (stop_found, text_before_stop)
                - stop_found: True if stop sequence found
                - text_before_stop: Text before the stop sequence (or full text if not found)
        """
        if not stop_sequences:
            return False, accumulated_text

        for stop_seq in stop_sequences:
            if stop_seq in accumulated_text:
                # Found stop sequence - return text before it
                idx = accumulated_text.index(stop_seq)
                return True, accumulated_text[:idx]

        return False, accumulated_text

    def _count_tokens_approx(self, text: str) -> int:
        """
        Approximate token count (4 chars per token).

        Args:
            text: Text to count

        Returns:
            int: Approximate token count
        """
        return len(text) // 4

    def _check_max_tokens(
        self,
        current_tokens: int,
        max_tokens: int | None
    ) -> bool:
        """
        Check if max tokens limit reached.

        Args:
            current_tokens: Current token count
            max_tokens: Maximum allowed tokens (None for unlimited)

        Returns:
            bool: True if limit reached
        """
        if max_tokens is None:
            return False
        return current_tokens >= max_tokens