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

Provides real token-by-token streaming for the REPL:
- True streaming from OpenAI/Anthropic APIs
- Token callbacks for real-time display
- Fallback to non-streaming for unsupported backends
- 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 os
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Optional

from groknroll.core.agent import GroknrollAgent
from groknroll.prompts.system import build_system_prompt, get_provider_prompt


@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 with TRUE streaming support.

    Uses the underlying LLM client's streaming capabilities for
    real-time token-by-token output.

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

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

        result = handler.complete_streaming_sync(
            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",
    ):
        """
        Initialize streaming handler

        Args:
            project_path: Project root path
            agent: GroknrollAgent instance (creates default if not provided)
            model: Default model to use
        """
        self.project_path = project_path or Path.cwd()
        self._agent = agent
        self.model = model
        self._client = None
        self._client_type = None

    @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
            config = AgentConfig(auto_index=False)
            self._agent = GroknrollAgent(self.project_path, config=config)
        return self._agent

    def _get_streaming_client(self):
        """Get or create a streaming-capable LLM client"""
        if self._client is not None:
            return self._client, self._client_type

        # Try Anthropic first (usually faster)
        anthropic_key = os.environ.get("ANTHROPIC_API_KEY")
        if anthropic_key:
            from groknroll.clients.anthropic import AnthropicClient
            self._client = AnthropicClient.get_instance(
                api_key=anthropic_key,
                model_name=self.model if "claude" in self.model else "claude-3-5-sonnet-20241022",
            )
            self._client_type = "anthropic"
            return self._client, self._client_type

        # Fall back to OpenAI
        openai_key = os.environ.get("OPENAI_API_KEY")
        if openai_key:
            from groknroll.clients.openai import OpenAIClient
            self._client = OpenAIClient.get_instance(
                api_key=openai_key,
                model_name=self.model,
            )
            self._client_type = "openai"
            return self._client, self._client_type

        return None, None

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

        Args:
            prompt: User prompt
            context: Additional context (file contents, etc.)
            on_token: Callback for each token (called immediately as tokens arrive)
            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 streaming client
            client, client_type = self._get_streaming_client()

            if client is not None and on_token is not None:
                # TRUE STREAMING - tokens arrive as they're generated
                response_parts = []

                for chunk in client.completion_stream(full_prompt, model=used_model):
                    response_parts.append(chunk)
                    on_token(chunk)  # Immediate callback - no delay!

                response = "".join(response_parts)

                # Get actual usage if available
                input_tokens = getattr(client, 'last_prompt_tokens', 0) or len(full_prompt.split()) * 1.3
                output_tokens = getattr(client, 'last_completion_tokens', 0) or len(response.split()) * 1.3

            else:
                # Fallback: non-streaming (still no artificial delays)
                response = self.agent.chat(full_prompt)

                if on_token:
                    on_token(response)  # Send all at once

                input_tokens = len(full_prompt.split()) * 1.3
                output_tokens = len(response.split()) * 1.3

            execution_time = time.time() - start_time
            total_tokens = int(input_tokens + output_tokens)

            # Estimate cost
            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),
            )

    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 TRUE streaming (async version)

        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:
            full_prompt = self._build_prompt(prompt, context)
            client, client_type = self._get_streaming_client()

            if client is not None and on_token is not None:
                # TRUE ASYNC STREAMING
                response_parts = []

                async for chunk in client.acompletion_stream(full_prompt, model=used_model):
                    response_parts.append(chunk)
                    on_token(chunk)

                response = "".join(response_parts)
                input_tokens = getattr(client, 'last_prompt_tokens', 0) or len(full_prompt.split()) * 1.3
                output_tokens = getattr(client, 'last_completion_tokens', 0) or len(response.split()) * 1.3

            else:
                response = self.agent.chat(full_prompt)
                if on_token:
                    on_token(response)
                input_tokens = len(full_prompt.split()) * 1.3
                output_tokens = len(response.split()) * 1.3

            execution_time = time.time() - start_time
            total_tokens = int(input_tokens + output_tokens)
            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 _get_system_prompt(self) -> str:
        """Get the comprehensive system prompt based on current model/provider."""
        # Determine provider from model name or client type
        provider_id = self._client_type or "anthropic"
        model_id = self.model or ""

        # Build system prompt
        prompts = build_system_prompt(
            provider_id=provider_id,
            model_id=model_id,
            agent_name="build",
            project_path=self.project_path,
            include_environment=True,
        )

        return "\n\n".join(prompts)

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

        # Add system prompt for better agent behavior
        system_prompt = self._get_system_prompt()
        if system_prompt:
            parts.append(system_prompt)

        # Add file contents from context
        if context and "files" in context:
            file_parts = []
            for file_path, content in context["files"].items():
                file_parts.append(f"File: {file_path}\n```\n{content}\n```")
            if file_parts:
                parts.append("# Referenced Files\n" + "\n\n".join(file_parts))

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

        # Add the user task
        parts.append(f"# Task\n{prompt}")

        return "\n\n".join(parts)

    def _estimate_cost(
        self,
        input_tokens: int,
        output_tokens: int,
        model: str,
    ) -> float:
        """Estimate cost based on model pricing"""
        # Pricing per 1M tokens
        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-20241022": {"input": 3.00, "output": 15.00},
            "claude-3-5-sonnet": {"input": 3.00, "output": 15.00},
            "claude-3-haiku": {"input": 0.25, "output": 1.25},
        }

        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"""
        return await self.complete_streaming(prompt, context, on_token)

    def set_model(self, model: str) -> None:
        """Set the default model"""
        self.model = model
        # Reset client to pick up new model
        self._client = None
        self._client_type = None


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