# Copyright 2025-2026 Horizon RL Contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""SGLang model provider with token-in/token-out support."""

from __future__ import annotations

import base64
import json
import logging
from collections.abc import AsyncGenerator, AsyncIterable, Callable, Iterator
from typing import (
    TYPE_CHECKING,
    Any,
    TypedDict,
    TypeVar,
    cast,
)

from pydantic import BaseModel
from strands.models import Model
from strands.types.content import ContentBlock, Messages, SystemContentBlock
from strands.types.exceptions import (
    ContextWindowOverflowException,
    ModelThrottledException,
)
from strands.types.streaming import StopReason, StreamEvent
from strands.types.tools import ToolChoice, ToolResultContent, ToolSpec
from typing_extensions import Unpack, override

from .client import SGLangClient
from .exceptions import SGLangContextLengthError, SGLangThrottledError
from .token import TokenManager
from .tool_parsers import HermesToolParser, ToolParser, ToolParseResult

if TYPE_CHECKING:
    from transformers import PreTrainedTokenizerBase, ProcessorMixin

logger = logging.getLogger(__name__)

T = TypeVar("T", bound=BaseModel)


class SGLangModel(Model):
    """SGLang native `/generate` API provider with token-in/token-out support.

    Example:
        >>> from transformers import AutoTokenizer
        >>> from strands_sglang import SGLangClient, SGLangModel
        >>> client = SGLangClient(base_url="http://localhost:30000")
        >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")
        >>> model = SGLangModel(client=client, tokenizer=tokenizer)
        >>> # After generation:
        >>> model.token_manager.token_ids    # Full token trajectory
        >>> model.token_manager.loss_mask    # Boolean mask for loss computation
        >>> model.token_manager.logprobs     # Log probabilities
    """

    class SGLangConfig(TypedDict, total=False):
        """Configuration options for SGLang generation."""

        sampling_params: dict[str, Any] | None  # Passed to /generate endpoint
        return_logprob: bool | None  # Return logprobs for all tokens (default: True)
        enable_thinking: bool | None  # Enable thinking mode for Qwen3 hybrid models

    def __init__(
        self,
        *,
        client: SGLangClient,
        tokenizer: PreTrainedTokenizerBase | None = None,
        processor: ProcessorMixin | None = None,
        tool_parser: ToolParser | None = None,
        **config: Unpack[SGLangConfig],
    ) -> None:
        """Initialize SGLang model provider.

        Args:
            client: `SGLangClient` for HTTP communication with the SGLang server.
            tokenizer: HuggingFace tokenizer for chat template and tokenization (optional if processor is provided).
            processor: HuggingFace processor for multimodal processing.
            tool_parser: `ToolParser` for tool calls (default: `HermesToolParser`).
            **config: Additional SGLang generation configuration.
        """
        self.client = client
        self.processor = processor
        self.tokenizer = cast(
            "PreTrainedTokenizerBase", (processor and getattr(processor, "tokenizer", None)) or tokenizer
        )
        if not self.tokenizer:
            raise ValueError("Either tokenizer (text-only) or processor (multimodal) must be provided")
        self.tool_parser = tool_parser or HermesToolParser()
        self.config = dict(config)
        self.tool_parser.validate_tokenizer(self.tokenizer)

        # State tracking (this makes SGLangModel stateful)
        self.token_manager = TokenManager()
        self._processed_message_count: int = 0
        self.tool_parse_errors: dict[str, int] = {}  # per-tool parse error count
        self.image_data: list[str] = []  # accumulated image data URLs (VLM only)

        logger.debug("initialized with config: %s", self.config)

    def reset(self) -> None:
        """Reset token accumulation for a new episode.

        Call this at episode start. Clears all accumulated tokens and resets
        internal state for tool tracking.
        """
        self.token_manager.reset()
        self._processed_message_count = 0
        self.tool_parse_errors = {}
        self.image_data = []

    @property
    def is_multimodal(self) -> bool:
        """Whether the model is multimodal."""
        return self.processor is not None

    # -------------------------------------------------------------------------
    # Model interface implementation
    # -------------------------------------------------------------------------

    @override
    def update_config(self, **model_config: Unpack[SGLangConfig]) -> None:  # type: ignore[override]
        """Update the model configuration.

        Args:
            **model_config: Configuration overrides.
        """
        self.config.update(model_config)

    @override
    def get_config(self) -> SGLangConfig:
        """Get the model configuration.

        Returns:
            The model configuration dict.
        """
        return cast(SGLangModel.SGLangConfig, self.config)

    # -------------------------------------------------------------------------
    # Chat template and message formatting
    # -------------------------------------------------------------------------

    @classmethod
    def format_content_block(
        cls, content_block: ContentBlock | ToolResultContent, is_multimodal: bool = False
    ) -> dict[str, Any] | str:
        """Convert a single Strands `ContentBlock` or `ToolResultContent` to HF chat template format."""
        # ContentBlock / ToolResultContent is a TypedDict with exactly one key set at runtime
        ((key, value),) = content_block.items()
        result: dict[str, Any] = {}
        match key, value:
            case "text", str() as text:
                result = {"type": "text", "text": text}
            case "image", dict() as image:
                mime = f"image/{image['format']}"
                encoded = base64.b64encode(image["source"]["bytes"]).decode()
                result = {"type": "image", "image": f"data:{mime};base64,{encoded}"}
            case "json", data:
                result = {"type": "text", "text": json.dumps(data)}
            # TODO: add support for other content types
            case _:
                raise TypeError(f"content_type=<{key}> | unsupported type")
        # flatten to text if not multimodal
        if not is_multimodal:
            return str(result["text"])
        return result

    @classmethod
    def format_messages(
        cls, messages: Messages, system_prompt: str | None = None, is_multimodal: bool = False
    ) -> list[dict[str, Any]]:
        """Convert Strands Messages to HF chat template format.

        When `is_multimodal=False` (default), content is flattened to a plain string.
        When `is_multimodal=True`, content is kept as a list of dicts.
        """
        result: list[dict[str, Any]] = []

        if system_prompt:
            result.append({"role": "system", "content": system_prompt})

        # Each Strands message is {"role": str, "content": [ContentBlock, ...]}
        # One Strands message maps to one HF message, except toolResult blocks
        # which each become a separate HF message with role="tool".
        for msg in messages:
            if "toolResult" in msg["content"][0]:
                # Each toolResult → its own HF message (different tool_call_id)
                for cb in msg["content"]:
                    assert "toolResult" in cb
                    tr = cb["toolResult"]
                    content = [cls.format_content_block(c, is_multimodal) for c in tr["content"]]
                    result.append(
                        {
                            "role": "tool",
                            "tool_call_id": tr["toolUseId"],
                            "content": content if is_multimodal else content[0],
                        }
                    )
            else:
                # Non-tool content → one HF message (text, image, etc.; toolUse skipped)
                content = [cls.format_content_block(c, is_multimodal) for c in msg["content"] if "toolUse" not in c]
                result.append({"role": msg["role"], "content": content if is_multimodal else content[0]})

        return result

    def format_tool_specs(self, tool_specs: list[ToolSpec]) -> list[dict]:
        """Format strands ToolSpecs to OpenAI format for chat templates."""
        return [
            {
                "type": "function",
                "function": {
                    "name": spec["name"],
                    "description": spec["description"],
                    "parameters": spec["inputSchema"]["json"],
                },
            }
            for spec in tool_specs
        ]

    def format_prompt(
        self,
        messages: Messages,
        system_prompt: str | None = None,
        tools: list[dict] | None = None,
        add_generation_prompt: bool = True,
    ) -> str:
        """Format messages into a prompt string using the HuggingFace chat template.

        The result is manually tokenized (not model-generated) and added to
        the token trajectory with `loss_mask=False`.
        """
        chat_messages = self.format_messages(messages, system_prompt, is_multimodal=self.is_multimodal)
        self.image_data.extend(self.extract_image_urls(chat_messages))
        # TODO: add support for other modalities later
        return str(
            self.tokenizer.apply_chat_template(
                conversation=chat_messages,
                tools=cast(list[dict | Callable], tools),
                add_generation_prompt=add_generation_prompt,
                tokenize=False,
                enable_thinking=self.config.get("enable_thinking"),
            )
        )

    @staticmethod
    def extract_image_urls(messages: list[dict[str, Any]]) -> list[str]:
        """Extract image data URLs from HF-formatted multimodal messages."""
        urls: list[str] = []
        for msg in messages:
            content = msg.get("content")
            if isinstance(content, dict) and content.get("type") == "image":
                urls.append(content["image"])
            elif isinstance(content, list):
                for part in content:
                    if isinstance(part, dict) and part.get("type") == "image":
                        urls.append(part["image"])
        return urls

    # -------------------------------------------------------------------------
    # Generation
    # -------------------------------------------------------------------------

    def tokenize_prompt_messages(
        self,
        messages: Messages,
        system_prompt: str | None,
        tools: list[dict] | None = None,
    ) -> list[int] | None:
        """Tokenize prompt messages for the next generation call.

        First call: tokenizes full prompt with system prompt and tools.
        Subsequent calls: tokenizes only new messages (tool results, user messages),
        prepending the message separator to align with chat template formatting.

        For VLM (when `self.processor` is set), uses the processor to insert
        image placeholder tokens based on `self.image_data`.
        """

        def _tokenize(text: str) -> list[int]:
            if self.processor:
                return list(self.processor(text=text, images=self.image_data or None)["input_ids"][0])  # type: ignore[arg-type]
            return list(self.tokenizer.encode(text, add_special_tokens=False))

        # First call: full prompt with tools
        if len(self.token_manager) == 0:
            formatted = self.format_prompt(messages, system_prompt, tools=tools)
            return _tokenize(formatted)

        # Subsequent calls: only new messages (prepend fake user to satisfy chat template assumptions).
        # See: https://github.com/horizon-rl/strands-sglang/issues/29
        if len(messages) > self._processed_message_count:
            new_messages = self._sort_tool_results(messages[self._processed_message_count :])
            fake: Messages = [{"role": "user", "content": [{"text": "ONLY FOR INCREMENTAL TOKENIZATION"}]}]
            full = self.format_prompt(fake + new_messages)
            prefix = self.format_prompt(fake, add_generation_prompt=False)
            formatted = self.tool_parser.message_separator + full[len(prefix) :]
            return _tokenize(formatted)

        return None

    def _sort_tool_results(self, messages: Messages) -> Messages:
        """Sort tool results by ID to match original call order (IDs are sequential: call_0000, call_0001, ...)."""
        result = []
        for msg in messages:
            if msg.get("role") != "user" or not isinstance(msg.get("content"), list):
                result.append(msg)
                continue
            content = msg["content"]
            tool_results = [b for b in content if isinstance(b, dict) and "toolResult" in b]
            if not tool_results:
                result.append(msg)
                continue
            other = [b for b in content if not (isinstance(b, dict) and "toolResult" in b)]
            tool_results.sort(key=lambda b: b.get("toolResult", {}).get("toolUseId", ""))
            result.append({**msg, "content": other + tool_results})
        return result

    def _yield_tool_use_events(
        self,
        tool_calls: list[ToolParseResult],
    ) -> Iterator[StreamEvent]:
        """Yield toolUse stream events for parsed tool calls.

        Each tool call emits three events following the Strands streaming protocol:
        - `contentBlockStart`: begins block with toolUseId and name
        - `contentBlockDelta`: contains the tool input (delta = incremental data)
        - `contentBlockStop`: ends the block
        """
        for tool_call in tool_calls:
            if tool_call.is_error:
                logger.warning("Tool parse error for '%s': %s", tool_call.name, (tool_call.raw or "")[:100])
                # Track parse error count per tool name
                self.tool_parse_errors[tool_call.name] = self.tool_parse_errors.get(tool_call.name, 0) + 1

            yield {
                "contentBlockStart": {
                    "start": {
                        "toolUse": {
                            "toolUseId": tool_call.id,
                            "name": tool_call.name,
                        }
                    }
                }
            }
            yield {
                "contentBlockDelta": {
                    "delta": {
                        "toolUse": {
                            "input": tool_call.payload,
                        }
                    }
                }
            }
            yield {"contentBlockStop": {}}

    def _extract_logprobs(self, event: dict[str, Any], key: str) -> list[float] | None:
        """Extract logprobs from SGLang event (format: [[logprob, token_id, ...], ...])."""
        meta_info = event.get("meta_info", {})
        logprobs = meta_info.get(key) or event.get(key)
        if isinstance(logprobs, list) and logprobs:
            return [entry[0] for entry in logprobs]
        return None

    @override
    async def stream(
        self,
        messages: Messages,
        tool_specs: list[ToolSpec] | None = None,
        system_prompt: str | None = None,
        *,
        tool_choice: ToolChoice | None = None,
        system_prompt_content: list[SystemContentBlock] | None = None,
        **kwargs: Any,
    ) -> AsyncIterable[StreamEvent]:
        """Chat completion with SGLangModel using the `/generate` endpoint.

        The `stream` method follows Strands' protocol but actually disabled here for training-only usage.
        This means users won't see streaming behavior such as print callbacks.
        """
        # Prepare request
        tools = self.format_tool_specs(tool_specs) if tool_specs else None
        config = self.get_config()
        sampling_params: dict[str, Any] = dict(config.get("sampling_params") or {})
        sampling_params.setdefault("skip_special_tokens", False)
        return_logprob = config.get("return_logprob", True)
        new_input_tokens = self.tokenize_prompt_messages(messages, system_prompt, tools=tools)
        # Tracking token IDs in token_manager to ensure the token-in feature
        input_ids = self.token_manager.token_ids + (new_input_tokens or [])

        # Start message
        yield {"messageStart": {"role": "assistant"}}
        yield {"contentBlockStart": {"start": {}}}

        # Call SGLangClient (non-streaming POST for better parallelism)
        try:
            response = await self.client.generate(
                input_ids=input_ids,
                sampling_params=sampling_params,
                return_logprob=return_logprob,
                logprob_start_len=0 if return_logprob else None,
                image_data=self.image_data or None,
            )

            # Extract response data
            text = response.get("text", "")
            output_ids = response.get("output_ids", [])
            output_logprobs = self._extract_logprobs(response, "output_token_logprobs")
            input_logprobs = self._extract_logprobs(response, "input_token_logprobs")
            meta_info = response.get("meta_info", {})

            # Yield text as single delta (non-streaming gives complete text at once)
            if text:
                yield {"contentBlockDelta": {"delta": {"text": text}}}

        except SGLangContextLengthError as e:
            raise ContextWindowOverflowException(f"Context length exceeded: {e.body}") from e
        except SGLangThrottledError as e:
            raise ModelThrottledException(f"Service throttled (status={e.status}): {e.body}") from e

        # Update token trajectory
        if new_input_tokens:
            new_input_logprobs = input_logprobs[-len(new_input_tokens) :] if input_logprobs else None
            self.token_manager.add_prompt(token_ids=new_input_tokens, logprobs=new_input_logprobs)
        if output_ids:
            self.token_manager.add_response(token_ids=output_ids, logprobs=output_logprobs)
        self._processed_message_count = len(messages) + 1

        # End text block, start tool use blocks if there are any tool calls
        yield {"contentBlockStop": {}}

        # Parse tool calls and yield events
        parsed_tool_calls = self.tool_parser.parse(text)
        for event in self._yield_tool_use_events(parsed_tool_calls):
            yield event

        # Determine stop reason
        stop_reason: str = "tool_use" if parsed_tool_calls else "end_turn"
        if meta_info and isinstance(meta_info.get("finish_reason"), dict):
            if meta_info["finish_reason"].get("type") == "length":
                stop_reason = "max_tokens"

        yield {"messageStop": {"stopReason": cast(StopReason, stop_reason)}}

        # Yield usage metadata
        if meta_info:
            prompt_tokens = int(meta_info.get("prompt_tokens") or 0)
            completion_tokens = int(meta_info.get("completion_tokens") or 0)
            yield {
                "metadata": {
                    "usage": {
                        "inputTokens": prompt_tokens,
                        "outputTokens": completion_tokens,
                        "totalTokens": prompt_tokens + completion_tokens,
                    },
                    "metrics": {"latencyMs": int(float(meta_info.get("e2e_latency") or 0) * 1000)},
                }
            }

    @override
    async def structured_output(
        self,
        output_model: type[T],
        prompt: Messages,
        system_prompt: str | None = None,
        **kwargs: Any,
    ) -> AsyncGenerator[dict[str, T | Any], None]:
        """Get structured output using SGLang's constrained decoding.

        Uses SGLang's `json_schema` parameter for FSM-based constrained generation,
        guaranteeing output conforms to the Pydantic model schema.

        Note: This method does NOT update token_manager (no TITO tracking).
        Intended for inference-only use cases like LLM-as-Judge.

        Args:
            output_model: Pydantic model class defining the output schema.
            prompt: Messages to send to the model.
            system_prompt: Optional system prompt.
            **kwargs: Additional arguments (unused).

        Yields:
            Single dict with "output" key containing the parsed Pydantic model instance.

        Raises:
            ValidationError: If model output fails Pydantic validation.
            SGLangHTTPError: On non-retryable HTTP errors.
        """
        # Convert Pydantic model to JSON schema string
        json_schema = json.dumps(output_model.model_json_schema())

        # Format and tokenize prompt (no tools for structured output)
        formatted = self.format_prompt(prompt, system_prompt, tools=None)
        input_ids = self.tokenizer.encode(formatted, add_special_tokens=False)

        # Build sampling params with json_schema constraint
        config = self.get_config()
        sampling_params: dict[str, Any] = dict(config.get("sampling_params") or {})
        sampling_params["json_schema"] = json_schema

        # Call SGLang /generate endpoint
        try:
            response = await self.client.generate(
                input_ids=input_ids,
                sampling_params=sampling_params,
                return_logprob=False,  # No need for logprobs in structured output
            )
        except SGLangContextLengthError as e:
            raise ContextWindowOverflowException(f"Context length exceeded: {e.body}") from e
        except SGLangThrottledError as e:
            raise ModelThrottledException(f"Service throttled (status={e.status}): {e.body}") from e

        # Parse and validate response
        text = response.get("text", "")
        parsed = output_model.model_validate_json(text)

        yield {"output": parsed}
