"""Stream Context - Manages streaming state and content accumulation

This module follows the Composition over Inheritance principle initially,
but StreamContext inherits from StreamState for simplicity.
The rendering logic is delegated to a separate StreamRenderer.
"""
from dataclasses import dataclass
from typing import List, Dict, Any, Optional
from enum import Enum

from luxx.services.events import sse_event


class StepType(str, Enum):
    """Step type enumeration"""
    THINKING = "thinking"
    TEXT = "text"
    TOOL_CALL = "tool_call"
    TOOL_RESULT = "tool_result"


THINK_START = "<think>"
THINK_END = "</think>"


@dataclass
class Step:
    """Represents a single step in the response process"""
    id: str
    index: int
    type: str
    content: str = ""
    name: str = ""
    arguments: str = ""
    id_ref: str = ""
    success: bool = True
    
    def to_dict(self) -> Dict[str, Any]:
        result = {
            "id": self.id,
            "index": self.index,
            "type": self.type,
        }
        if self.content:
            result["content"] = self.content
        if self.name:
            result["name"] = self.name
        if self.arguments:
            result["arguments"] = self.arguments
        if self.id_ref:
            result["id_ref"] = self.id_ref
        if self.type == StepType.TOOL_RESULT:
            result["success"] = self.success
        return result


class StreamState:
    """Pure state management for streaming
    
    This class maintains all state but delegates rendering to StreamRenderer.
    """
    
    def __init__(self):
        self.reset()
    
    def reset(self, full_reset: bool = True):
        """Reset state for a new iteration or full reset.
        
        Args:
            full_reset: If True, reset everything. If False, only reset per-iteration state.
        """
        if full_reset:
            # Full reset - clear everything
            self.step_index = 0
            self.all_steps: List[Step] = []
            self.all_tool_results: List[Dict] = []
        
        # Per-iteration reset (always runs)
        self.current_step_id: Optional[str] = None
        self.current_step_idx: Optional[int] = None
        self.current_step_type: Optional[str] = None
        self.full_content = ""
        self.full_thinking = ""
        self.tool_calls_list: List[Dict] = []
        self._last_message_id: Optional[str] = None
        self._last_token_count = 0
        self._last_usage: Optional[Dict] = None
        self._in_thinking = False
        self._thinking_buf = ""
        self._text_buf = ""
        # Track content offset per step to avoid duplication across steps
        self._text_offset = 0
        self._thinking_offset = 0
    
    def process_content(self, content: str) -> Dict:
        """Process raw content, handling thinking tags."""
        if not content:
            return {"thinking": "", "text": "", "should_emit": False, "thinking_only": False}
        
        thinking = ""
        text = ""
        should_emit = False
        thinking_only = False
        
        # Handle THINK_START - can appear anywhere in content
        if not self._in_thinking and THINK_START in content:
            self._in_thinking = True
            idx = content.find(THINK_START)
            # Any text before THINK_START is emitted immediately
            if idx > 0:
                text += content[:idx]
            content = content[idx + len(THINK_START):]
        
        # Handle THINK_END - can appear anywhere
        if THINK_END in content:
            idx = content.find(THINK_END)
            thinking_content = content[:idx]
            content = content[idx + len(THINK_END):]
            
            self._in_thinking = False
            should_emit = True
            thinking = thinking_content
            
            # Handle any remaining text after THINK_END
            while THINK_END in content:
                second_idx = content.find(THINK_END)
                text += content[:second_idx]
                content = content[second_idx + len(THINK_END):]
            
            # Any remaining content after last THINK_END is text
            if content:
                text += content
            
            thinking_only = not bool(text)
        elif self._in_thinking:
            # In thinking mode, stream content directly
            thinking = content
            should_emit = True
            thinking_only = True
        else:
            # Not in thinking mode, emit text immediately for streaming
            text = content
            should_emit = True
        
        # Emit text that appeared before THINK_START
        if text and not should_emit:
            should_emit = True
        
        return {
            "thinking": thinking,
            "text": text,
            "should_emit": should_emit,
            "thinking_only": thinking_only
        }
    
    def flush(self) -> tuple:
        """Flush remaining buffers and return content"""
        thinking = self._thinking_buf
        text = self._text_buf
        self._thinking_buf = ""
        self._text_buf = ""
        return thinking, text
    
    def start_step(self, step_type: str) -> str:
        """Start a new step and return its ID. Auto-finalizes previous step."""
        # Auto-finalize previous step before starting new one
        self.finalize_step()
        
        self.current_step_idx = self.step_index
        self.current_step_id = f"step-{self.step_index}"
        self.current_step_type = step_type
        # Record content offset so this step only includes content added during it
        if step_type == StepType.TEXT:
            self._text_offset = len(self.full_content)
        elif step_type == StepType.THINKING:
            self._thinking_offset = len(self.full_thinking)
        self.step_index += 1
        return self.current_step_id
    
    def finalize_step(self):
        """Finalize the current step and add to all_steps"""
        if self.current_step_id is None:
            return
        if self.current_step_type == StepType.TEXT:
            content = self.full_content[self._text_offset:]
        elif self.current_step_type == StepType.THINKING:
            content = self.full_thinking[self._thinking_offset:]
        else:
            content = ""
        step = Step(
            id=self.current_step_id,
            index=self.current_step_idx,
            type=self.current_step_type,
            content=content
        )
        self.all_steps.append(step)
        # Reset current step state but keep buffers for accumulation
        self.current_step_id = None
        self.current_step_idx = None
        self.current_step_type = None
    
    def accumulate_tool_call(self, tc_delta: Dict):
        """Accumulate tool call delta"""
        idx = tc_delta.get("index", 0)
        if idx >= len(self.tool_calls_list):
            self.tool_calls_list.append({
                "id": tc_delta.get("id", ""),
                "type": "function",
                "function": {"name": "", "arguments": ""}
            })
        func = tc_delta.get("function", {})
        if func.get("name"):
            self.tool_calls_list[idx]["function"]["name"] += func["name"]
        if func.get("arguments"):
            self.tool_calls_list[idx]["function"]["arguments"] += func["arguments"]
    
    def add_tool_result(self, result: Dict):
        """Add a tool result to history"""
        self.all_tool_results.append({
            "role": "tool",
            "tool_call_id": result.get("tool_call_id", ""),
            "content": result.get("content", "")
        })
    
    def set_completion(self, msg_id: str, token_count: int, usage: dict):
        """Set completion metadata"""
        self._last_message_id = msg_id
        self._last_token_count = token_count
        self._last_usage = usage
    
    def get_steps_for_save(self) -> List[Dict]:
        """Get all steps as dictionaries"""
        return [step.to_dict() for step in self.all_steps]


class StreamRenderer:
    """Renders stream state to SSE events"""
    
    @staticmethod
    def render_tool_calls(state: StreamState) -> List[str]:
        """Render tool calls as SSE events"""
        events = []
        for tc in state.tool_calls_list:
            # Use start_step to auto-finalize previous and create new step
            state.start_step(StepType.TOOL_CALL)
            step = Step(
                id=state.current_step_id,
                index=state.current_step_idx,
                type=StepType.TOOL_CALL,
                name=tc["function"]["name"],
                arguments=tc["function"]["arguments"],
                id_ref=tc.get("id", "")
            )
            # Append again since start_step finalized previous (if any)
            state.all_steps.append(step)
            events.append(sse_event("process_step", {"step": step.to_dict()}))
        return events
    
    @staticmethod
    def render_tool_result(state: StreamState, result: Dict, ref_step_id: str) -> tuple:
        """Render a tool result as SSE event"""
        import json
        
        # Use start_step to auto-finalize previous and create new step
        state.start_step(StepType.TOOL_RESULT)
        content = result.get("content", "")
        success = True
        
        try:
            parsed = json.loads(content)
            if isinstance(parsed, dict):
                success = parsed.get("success", True)
        except (json.JSONDecodeError, TypeError):
            pass
        
        step = Step(
            id=state.current_step_id,
            index=state.current_step_idx,
            type=StepType.TOOL_RESULT,
            name=result.get("name", ""),
            content=content,
            id_ref=ref_step_id,
            success=success
        )
        state.all_steps.append(step)
        state.add_tool_result(result)
        
        return step, sse_event("process_step", {"step": step.to_dict()})
    
    @staticmethod
    def render_thinking(state: StreamState) -> str:
        """Render thinking content as SSE event"""
        step = Step(
            id=state.current_step_id,
            index=state.current_step_idx,
            type=StepType.THINKING,
            content=state.full_thinking[state._thinking_offset:]
        )
        return sse_event("process_step", {"step": step.to_dict()})
    
    @staticmethod
    def render_text(state: StreamState) -> str:
        """Render text content as SSE event"""
        step = Step(
            id=state.current_step_id,
            index=state.current_step_idx,
            type=StepType.TEXT,
            content=state.full_content[state._text_offset:]
        )
        return sse_event("process_step", {"step": step.to_dict()})
    
    @staticmethod
    def render_error(error_msg: str) -> str:
        """Render error event"""
        return sse_event("error", {"content": error_msg})


# Convenience function for backward compatibility
def render_error(error_msg: str) -> str:
    """Render error event"""
    return sse_event("error", {"content": error_msg})