"""
Streaming RLM Handler - Streaming responses from the LLM

Provides streaming completion for the REPL:
- Token-by-token output with callbacks
- Fallback chunking for non-streaming backends
- Integration with RLM core
- Cost and timing tracking

Examples:
    handler = StreamingRLMHandler(project_path=Path.cwd())

    # Stream with callback
    def on_token(token):
        print(token, end="", flush=True)

    result = await handler.complete_streaming(
        prompt="Explain this code",
        context={"file": "main.py"},
        on_token=on_token,
    )
    print(f"\nCost: ${result.cost:.4f}")
"""

import asyncio
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Optional

from groknroll.core.agent import GroknrollAgent


@dataclass
class RLMResult:
    """
    Result from RLM completion

    Attributes:
        response: Complete response text
        cost: Cost in dollars
        tokens: Total tokens used
        input_tokens: Input tokens
        output_tokens: Output tokens
        execution_time: Time taken in seconds
        model: Model used
        success: Whether completion succeeded
        error: Error message if failed
    """

    response: str
    cost: float = 0.0
    tokens: int = 0
    input_tokens: int = 0
    output_tokens: int = 0
    execution_time: float = 0.0
    model: str = ""
    success: bool = True
    error: Optional[str] = None

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary"""
        return {
            "response": self.response,
            "cost": self.cost,
            "tokens": self.tokens,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "execution_time": self.execution_time,
            "model": self.model,
            "success": self.success,
            "error": self.error,
        }


class StreamingRLMHandler:
    """
    Handle streaming LLM completions

    Provides streaming responses with token callbacks for real-time
    display in the REPL.

    Example:
        handler = StreamingRLMHandler(project_path=Path.cwd())

        async def display_token(token: str) -> None:
            print(token, end="", flush=True)

        result = await handler.complete_streaming(
            prompt="Fix the bug",
            context={},
            on_token=display_token,
        )
    """

    def __init__(
        self,
        project_path: Optional[Path] = None,
        agent: Optional[GroknrollAgent] = None,
        model: str = "gpt-4o-mini",
        chunk_size: int = 10,
        chunk_delay: float = 0.02,
    ):
        """
        Initialize streaming handler

        Args:
            project_path: Project root path
            agent: GroknrollAgent instance (creates default if not provided)
            model: Default model to use
            chunk_size: Characters per chunk for fallback chunking
            chunk_delay: Delay between chunks in seconds
        """
        self.project_path = project_path or Path.cwd()
        self._agent = agent
        self.model = model
        self.chunk_size = chunk_size
        self.chunk_delay = chunk_delay

    @property
    def agent(self) -> GroknrollAgent:
        """Get or create the GroknrollAgent"""
        if self._agent is None:
            from groknroll.core.agent import AgentConfig

            # Disable auto_index for faster startup - index on demand instead
            config = AgentConfig(auto_index=False)
            self._agent = GroknrollAgent(self.project_path, config=config)
        return self._agent

    async def complete_streaming(
        self,
        prompt: str,
        context: Optional[dict] = None,
        on_token: Optional[Callable[[str], None]] = None,
        model: Optional[str] = None,
    ) -> RLMResult:
        """
        Complete with streaming token output

        Args:
            prompt: User prompt
            context: Additional context (file contents, etc.)
            on_token: Callback for each token
            model: Model to use (overrides default)

        Returns:
            RLMResult with complete response
        """
        start_time = time.time()
        used_model = model or self.model

        try:
            # Build full prompt with context
            full_prompt = self._build_prompt(prompt, context)

            # Get completion (RLM doesn't natively support streaming yet)
            # So we get the full response and chunk it
            response = self.agent.chat(full_prompt)

            # Simulate streaming by chunking the response
            if on_token:
                await self._stream_response(response, on_token)

            execution_time = time.time() - start_time

            # Estimate tokens (rough approximation)
            input_tokens = len(full_prompt.split()) * 1.3
            output_tokens = len(response.split()) * 1.3
            total_tokens = int(input_tokens + output_tokens)

            # Estimate cost (using gpt-4o-mini pricing as baseline)
            cost = self._estimate_cost(int(input_tokens), int(output_tokens), used_model)

            return RLMResult(
                response=response,
                cost=cost,
                tokens=total_tokens,
                input_tokens=int(input_tokens),
                output_tokens=int(output_tokens),
                execution_time=execution_time,
                model=used_model,
                success=True,
            )

        except Exception as e:
            execution_time = time.time() - start_time
            return RLMResult(
                response="",
                execution_time=execution_time,
                model=used_model,
                success=False,
                error=str(e),
            )

    def complete_streaming_sync(
        self,
        prompt: str,
        context: Optional[dict] = None,
        on_token: Optional[Callable[[str], None]] = None,
        model: Optional[str] = None,
    ) -> RLMResult:
        """
        Synchronous wrapper for complete_streaming

        Args:
            prompt: User prompt
            context: Additional context
            on_token: Token callback
            model: Model to use

        Returns:
            RLMResult
        """
        return asyncio.run(self.complete_streaming(prompt, context, on_token, model))

    async def _stream_response(
        self,
        response: str,
        on_token: Callable[[str], None],
    ) -> None:
        """
        Stream response in chunks

        Args:
            response: Full response text
            on_token: Callback for each chunk
        """
        for chunk in self._chunk_response(response):
            on_token(chunk)
            await asyncio.sleep(self.chunk_delay)

    def _chunk_response(self, response: str) -> list[str]:
        """
        Break response into chunks

        Args:
            response: Full response text

        Returns:
            List of chunks
        """
        chunks = []
        words = response.split(" ")
        current_chunk = []
        current_length = 0

        for word in words:
            if current_length + len(word) + 1 > self.chunk_size:
                if current_chunk:
                    chunks.append(" ".join(current_chunk) + " ")
                current_chunk = [word]
                current_length = len(word)
            else:
                current_chunk.append(word)
                current_length += len(word) + 1

        if current_chunk:
            chunks.append(" ".join(current_chunk))

        return chunks

    def _build_prompt(self, prompt: str, context: Optional[dict]) -> str:
        """
        Build full prompt with context

        Args:
            prompt: User prompt
            context: Additional context

        Returns:
            Full prompt string
        """
        if not context:
            return prompt

        context_parts = []

        # Add file contents
        if "files" in context:
            for file_path, content in context["files"].items():
                context_parts.append(f"File: {file_path}\n```\n{content}\n```")

        # Add other context
        for key, value in context.items():
            if key == "files":
                continue
            if isinstance(value, str):
                context_parts.append(f"{key}:\n{value}")

        if context_parts:
            context_str = "\n\n".join(context_parts)
            return f"Context:\n{context_str}\n\nTask: {prompt}"

        return prompt

    def _estimate_cost(
        self,
        input_tokens: int,
        output_tokens: int,
        model: str,
    ) -> float:
        """
        Estimate cost based on model pricing

        Args:
            input_tokens: Input token count
            output_tokens: Output token count
            model: Model name

        Returns:
            Estimated cost in dollars
        """
        # Pricing per 1M tokens (approximate)
        pricing = {
            "gpt-4o-mini": {"input": 0.15, "output": 0.60},
            "gpt-4o": {"input": 2.50, "output": 10.00},
            "gpt-4-turbo": {"input": 10.00, "output": 30.00},
            "claude-3-5-sonnet": {"input": 3.00, "output": 15.00},
            "claude-3-haiku": {"input": 0.25, "output": 1.25},
        }

        # Get pricing for model (default to gpt-4o-mini)
        model_pricing = pricing.get(model, pricing["gpt-4o-mini"])

        input_cost = (input_tokens / 1_000_000) * model_pricing["input"]
        output_cost = (output_tokens / 1_000_000) * model_pricing["output"]

        return input_cost + output_cost

    async def complete_with_tools(
        self,
        prompt: str,
        tools: list[dict],
        context: Optional[dict] = None,
        on_token: Optional[Callable[[str], None]] = None,
    ) -> RLMResult:
        """
        Complete with tool use support

        Args:
            prompt: User prompt
            tools: Available tools (OpenAI format)
            context: Additional context
            on_token: Token callback

        Returns:
            RLMResult (may include tool calls in response)
        """
        # For now, just do regular completion
        # Tool handling is done at a higher level
        return await self.complete_streaming(prompt, context, on_token)

    def set_model(self, model: str) -> None:
        """
        Set the default model

        Args:
            model: Model identifier
        """
        self.model = model


def create_streaming_handler(
    project_path: Optional[Path] = None,
    model: str = "gpt-4o-mini",
) -> StreamingRLMHandler:
    """
    Create a streaming handler (convenience function)

    Args:
        project_path: Project path
        model: Model to use

    Returns:
        Configured StreamingRLMHandler
    """
    return StreamingRLMHandler(project_path=project_path, model=model)
