import heapq
import logging
import math
from typing import Any, Dict, List, Optional, Union

from fastapi import Request
from fastapi.responses import ORJSONResponse

from sglang.srt.entrypoints.openai.protocol import (
    ChatCompletionMessageContentImagePart,
    ChatCompletionMessageContentTextPart,
    ChatCompletionMessageContentVideoPart,
    ErrorResponse,
    RerankContent,
    RerankResponse,
    V1RerankReqInput,
)
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput

logger = logging.getLogger(__name__)


def _get_yes_no_token_ids(tokenizer) -> tuple[int, int]:
    """Get token IDs for 'yes' and 'no' from the tokenizer.

    Different model sizes may have different token IDs, so we look them up dynamically.
    """
    # Try to encode 'yes' and 'no' to get their token IDs
    # The tokenizer should return a single token for these common words
    try:
        yes_tokens = tokenizer.encode("yes", add_special_tokens=False)
        no_tokens = tokenizer.encode("no", add_special_tokens=False)

        if len(yes_tokens) == 1 and len(no_tokens) == 1:
            return yes_tokens[0], no_tokens[0]

        # Fallback: try convert_tokens_to_ids
        yes_id = tokenizer.convert_tokens_to_ids("yes")
        no_id = tokenizer.convert_tokens_to_ids("no")
        if yes_id is not None and no_id is not None:
            return yes_id, no_id

    except Exception as e:
        logger.warning(f"Failed to get yes/no token IDs dynamically: {e}")

    # Fallback to known Qwen3 token IDs (may not work for all model sizes)
    logger.warning("Using fallback token IDs for yes/no (9693/2152)")
    return 9693, 2152


def _is_qwen3_reranker_template(chat_template: str) -> bool:
    """Detect if the chat template is for Qwen3 text-only reranker."""
    if not chat_template:
        return False
    t = chat_template.lower()
    return ('answer can only be "yes" or "no"' in t) or (
        "answer can only be" in t and '"yes"' in t and '"no"' in t
    )


def _is_qwen3_vl_reranker_template(chat_template: str) -> bool:
    """Detect if the chat template is for Qwen3-VL multimodal reranker.

    VL reranker templates use `query` and `document` as jinja variables
    and include vision token placeholders for image/video support.
    """
    if not chat_template:
        return False
    t = chat_template.lower()
    # Check for reranker phrase (yes/no judgment)
    has_reranker_phrase = ('answer can only be "yes" or "no"' in t) or (
        "answer can only be" in t and '"yes"' in t and '"no"' in t
    )
    # Check for vision token placeholders (unique to VL templates)
    has_vision_tokens = "<|vision_start|>" in t or "<|image_pad|>" in t
    return has_reranker_phrase and has_vision_tokens


def _is_qwen3_vl_model(model_path: str) -> bool:
    """Check if the model is a Qwen3-VL model based on model path."""
    if not model_path:
        return False
    model_lower = model_path.lower()
    return "qwen3-vl" in model_lower or "qwen3vl" in model_lower


def _detect_rerank_backend(
    *,
    request: V1RerankReqInput,
    chat_template: Optional[str],
    model_path: str,
) -> str:
    """
    Unify rerank routing decisions used by both `_convert_to_internal_request` and
    `_handle_non_streaming_request`.

    Returns:
        "vl_decoder" | "text_decoder" | "cross_encoder"
    """
    is_multimodal = request.is_multimodal()
    is_vl_model = _is_qwen3_vl_model(model_path)
    is_vl_template = _is_qwen3_vl_reranker_template(chat_template)
    is_text_template = _is_qwen3_reranker_template(chat_template)

    # Prefer VL when template/model indicates VL, or request is multimodal with reranker template.
    if is_vl_template or is_vl_model or (is_multimodal and is_text_template):
        return "vl_decoder"
    if is_text_template:
        return "text_decoder"
    return "cross_encoder"


def _qwen3_rerank_score(p_yes: float, p_no: float) -> float:
    denom = p_yes + p_no
    if denom <= 0.0:
        return 0.0
    return p_yes / denom


def _get_jinja_env():
    try:
        import jinja2  # Lazy import: server env should provide this dependency.
        from jinja2.sandbox import ImmutableSandboxedEnvironment
    except ModuleNotFoundError as e:
        raise ValueError(
            "Rendering Qwen3 reranker prompts requires `jinja2`. "
            "Please install it in your runtime environment (e.g., `pip install jinja2`)."
        ) from e
    # Using a sandboxed environment to stop malicious execution during model loading.
    return ImmutableSandboxedEnvironment(
        loader=jinja2.BaseLoader(),
        autoescape=False,
        undefined=jinja2.Undefined,
    )


def _render_jinja_chat_template(
    chat_template: str,
    *,
    query: RerankContent,
    document: RerankContent,
    instruct: Optional[str],
) -> str:
    """Render a loaded Jinja chat template for Qwen3 reranker prompts (text-only)."""
    env = _get_jinja_env()
    template = env.from_string(chat_template)

    # For text-only template, extract text content
    query_text = query if isinstance(query, str) else _extract_text_from_content(query)
    doc_text = (
        document if isinstance(document, str) else _extract_text_from_content(document)
    )

    render_kwargs = {
        "messages": [
            {"role": "user", "content": query_text},
            {"role": "user", "content": doc_text},
        ]
    }
    # Only pass instruct when explicitly provided; template uses `default(...)`
    # which works only when the variable is undefined (not None).
    if instruct:
        render_kwargs["instruct"] = instruct
    return template.render(**render_kwargs)


def _render_vl_jinja_template(
    chat_template: str,
    *,
    query: List[Dict[str, Any]],
    document: List[Dict[str, Any]],
    instruct: Optional[str],
) -> str:
    """Render a loaded Jinja chat template for Qwen3-VL reranker prompts (multimodal).

    The template expects `query` and `document` as lists of content parts,
    where each part has a `type` field (text, image, video) and corresponding data.
    """
    env = _get_jinja_env()
    template = env.from_string(chat_template)

    render_kwargs = {
        "query": query,
        "document": document,
    }
    if instruct:
        render_kwargs["instruct"] = instruct
    return template.render(**render_kwargs)


def _extract_text_from_content(content: RerankContent) -> str:
    """Extract text from multimodal content."""
    if isinstance(content, str):
        return content
    texts = []
    for part in content:
        if isinstance(part, ChatCompletionMessageContentTextPart):
            texts.append(part.text)
        elif isinstance(part, dict) and part.get("type") == "text":
            texts.append(part.get("text", ""))
    return " ".join(texts)


class OpenAIServingRerank(OpenAIServingBase):
    """Handler for /v1/rerank requests"""

    def __init__(self, tokenizer_manager, template_manager=None):
        super().__init__(tokenizer_manager)
        # TemplateManager is optional; rerank uses tokenizer.chat_template today.
        # Keeping this explicit makes the dependency clear and supports future extensions.
        self.template_manager = template_manager

        # Cache yes/no token IDs for Qwen3 reranker scoring
        self._yes_token_id, self._no_token_id = _get_yes_no_token_ids(
            tokenizer_manager.tokenizer
        )

    # NOTE: /v1/rerank is not an official OpenAI endpoint. This module may be moved
    # to another module in the future.

    def _request_id_prefix(self) -> str:
        return "rerank-"

    def _validate_request(self, request: V1RerankReqInput) -> Optional[str]:
        """Validate rerank request format and content"""
        if not request.query:
            return "Query cannot be empty"

        if isinstance(request.query, str):
            if not request.query.strip():
                return "Query cannot be empty or whitespace only"

        if not request.documents:
            return "Documents cannot be empty"

        for doc in request.documents:
            if not doc:
                return "Each document must be a non-empty string"
            if isinstance(doc, str) and not doc.strip():
                return "Each document cannot be empty or whitespace only"

        return None

    def _convert_to_internal_request(
        self,
        request: V1RerankReqInput,
        raw_request: Request = None,
    ) -> tuple[Union[EmbeddingReqInput, V1RerankReqInput], V1RerankReqInput]:
        """
        Convert OpenAI rerank request to internal format.

        - For Qwen3-VL reranker (multimodal decoder-only): keep the request.
        - For Qwen3 reranker (text-only decoder-only): keep the request and score via
          `tokenizer_manager.score_prompts(...)` in the handler.
        - For cross-encoder rerank models: adapt into `EmbeddingReqInput` pairs.
        """
        chat_template = self.tokenizer_manager.tokenizer.chat_template
        model_path = getattr(self.tokenizer_manager.model_config, "model_path", "")
        backend = _detect_rerank_backend(
            request=request,
            chat_template=chat_template if isinstance(chat_template, str) else None,
            model_path=model_path,
        )
        if backend in ("vl_decoder", "text_decoder"):
            return request, request

        # Cross-encoder rerank: Create pairs of [query, document] for each document.
        # Note: Cross-encoder only supports text-only content
        if request.is_multimodal():
            # Extract text for cross-encoder (multimodal not supported)
            query_text = _extract_text_from_content(request.query)
            doc_texts = [_extract_text_from_content(doc) for doc in request.documents]
            pairs = [[query_text, doc] for doc in doc_texts]
        else:
            pairs = [[request.query, doc] for doc in request.documents]

        adapted_request = EmbeddingReqInput(text=pairs, is_cross_encoder_request=True)
        return adapted_request, request

    async def _handle_non_streaming_request(
        self,
        adapted_request: Union[EmbeddingReqInput, V1RerankReqInput],
        request: V1RerankReqInput,
        raw_request: Request,
    ) -> Union[List[RerankResponse], ErrorResponse, ORJSONResponse]:
        """Handle the rerank request"""
        chat_template = getattr(self.tokenizer_manager.tokenizer, "chat_template", None)
        model_path = getattr(self.tokenizer_manager.model_config, "model_path", "")
        rerank_ret = await self._handle_rerank_paths(
            request=request,
            raw_request=raw_request,
            chat_template=chat_template,
            model_path=model_path,
        )
        if rerank_ret is not None:
            return rerank_ret

        # Default cross-encoder rerank path (existing behavior).
        try:
            if not isinstance(adapted_request, EmbeddingReqInput):
                raise ValueError(
                    "Invalid rerank request adaptation. "
                    "If you are serving a decoder-only reranker (e.g., Qwen3-Reranker), "
                    "please provide the corresponding --chat-template and launch without --is-embedding."
                )
            ret = await self.tokenizer_manager.generate_request(
                adapted_request, raw_request
            ).__anext__()
        except ValueError as e:
            return self.create_error_response(str(e))

        if not isinstance(ret, list):
            ret = [ret]

        responses = self._build_rerank_response(ret, request)
        return responses

    async def _handle_rerank_paths(
        self,
        *,
        request: V1RerankReqInput,
        raw_request: Request,
        chat_template: Optional[str],
        model_path: str,
    ) -> Optional[Union[List[RerankResponse], ErrorResponse, ORJSONResponse]]:
        """
        Handle decoder-only rerank paths (VL/text) and return a response if matched.

        Returns None if the request should fall back to cross-encoder rerank.
        """
        backend = _detect_rerank_backend(
            request=request,
            chat_template=chat_template,
            model_path=model_path,
        )

        # Qwen3-VL reranker path (decoder-only scoring with query/document template format)
        if backend == "vl_decoder":
            return await self._handle_vl_reranker_request(
                request, raw_request, chat_template or ""
            )

        # Qwen3 text-only reranker path (decoder-only scoring).
        if backend == "text_decoder":
            return await self._handle_text_reranker_request(
                request=request,
                raw_request=raw_request,
                chat_template=chat_template or "",
            )

        return None

    async def _handle_text_reranker_request(
        self,
        *,
        request: V1RerankReqInput,
        raw_request: Request,
        chat_template: str,
    ) -> Union[List[RerankResponse], ErrorResponse]:
        """Handle text-only decoder reranker request via score_prompts()."""
        # Qwen3 reranker relies on decoder-only logprobs. If the server is launched
        # with --is-embedding, model_config.is_generation is typically False and
        # logprob scoring is not supported.
        if not self.tokenizer_manager.model_config.is_generation:
            return self.create_error_response(
                "Detected Qwen3 reranker chat template, but the server is not in generation mode. "
                "Please relaunch without --is-embedding for Qwen3-Reranker models."
            )

        try:
            prompts = [
                _render_jinja_chat_template(
                    chat_template,
                    query=request.query,
                    document=doc,
                    instruct=getattr(request, "instruct", None),
                )
                for doc in request.documents
            ]

            result = await self.tokenizer_manager.score_prompts(
                prompts,
                label_token_ids=[self._yes_token_id, self._no_token_id],
                apply_softmax=False,
                request=raw_request,
            )
            scores = [_qwen3_rerank_score(s[0], s[1]) for s in result.scores]
        except ValueError as e:
            return self.create_error_response(str(e))
        except Exception as e:
            # Includes template rendering errors from jinja2.
            return self.create_error_response(str(e))

        responses = self._build_rerank_response(scores, request)
        return responses

    async def _handle_vl_reranker_request(
        self,
        request: V1RerankReqInput,
        raw_request: Request,
        _chat_template: str,
    ) -> Union[List[RerankResponse], ErrorResponse]:
        """Handle multimodal VL reranker request using chat completion with logprobs."""
        if not self.tokenizer_manager.model_config.is_generation:
            return self.create_error_response(
                "Detected Qwen3-VL reranker, but the server is not in generation mode. "
                "Please relaunch without --is-embedding for Qwen3-VL-Reranker models."
            )

        try:
            scores = []
            instruct = getattr(request, "instruct", None)

            for doc in request.documents:
                # Build multimodal content lists and render prompt using jinja template
                query_content, doc_content, image_data, video_data = (
                    self._build_vl_reranker_content(
                        query=request.query,
                        document=doc,
                    )
                )

                # Render the chat template directly with query/document variables
                prompt = _render_vl_jinja_template(
                    chat_template=_chat_template,
                    query=query_content,
                    document=doc_content,
                    instruct=instruct,
                )

                # Create generate request with logprobs
                gen_request = GenerateReqInput(
                    text=prompt,
                    image_data=image_data if image_data else None,
                    video_data=video_data if video_data else None,
                    sampling_params={
                        "max_new_tokens": 1,
                        "temperature": 0,
                    },
                    return_logprob=True,
                    top_logprobs_num=50,  # Get enough logprobs to find yes/no tokens
                    logprob_start_len=0,
                )

                # Execute generation request
                ret = await self.tokenizer_manager.generate_request(
                    gen_request, raw_request
                ).__anext__()

                # Extract yes/no probabilities from logprobs
                score = self._extract_score_from_logprobs(ret)
                scores.append(score)

            responses = self._build_rerank_response(scores, request)
            return responses

        except ValueError as e:
            return self.create_error_response(str(e))
        except Exception as e:
            logger.exception("Error handling VL reranker request")
            return self.create_error_response(str(e))

    def _build_vl_reranker_content(
        self,
        query: RerankContent,
        document: RerankContent,
    ) -> tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[str], List[str]]:
        """Build content lists for VL reranker request.

        Returns:
            Tuple of (query_content, document_content, image_data, video_data)
            where query_content and document_content are lists suitable for jinja template.
        """
        image_data = []
        video_data = []

        # Build query content list
        query_content = self._content_to_template_list(query, image_data, video_data)

        # Build document content list
        doc_content = self._content_to_template_list(document, image_data, video_data)

        return query_content, doc_content, image_data, video_data

    def _content_to_template_list(
        self,
        content: RerankContent,
        image_data: List[str],
        video_data: List[str],
    ) -> List[Dict[str, Any]]:
        """Convert RerankContent to a list format suitable for jinja template."""
        result = []

        if isinstance(content, str):
            result.append({"type": "text", "text": content})
            return result

        for part in content:
            if isinstance(part, ChatCompletionMessageContentTextPart):
                result.append({"type": "text", "text": part.text})
            elif isinstance(part, ChatCompletionMessageContentImagePart):
                if part.image_url:
                    image_data.append(part.image_url.url)
                    result.append({"type": "image"})
            elif isinstance(part, ChatCompletionMessageContentVideoPart):
                if part.video_url:
                    video_data.append(part.video_url.url)
                    result.append({"type": "video"})
            elif isinstance(part, dict):
                part_type = part.get("type")
                if part_type == "text":
                    result.append({"type": "text", "text": part.get("text", "")})
                elif part_type == "image_url":
                    image_url = part.get("image_url", {})
                    if isinstance(image_url, dict):
                        url = image_url.get("url")
                    else:
                        url = image_url
                    if url:
                        image_data.append(url)
                        result.append({"type": "image"})
                elif part_type == "video_url":
                    video_url = part.get("video_url", {})
                    if isinstance(video_url, dict):
                        url = video_url.get("url")
                    else:
                        url = video_url
                    if url:
                        video_data.append(url)
                        result.append({"type": "video"})

        return result

    def _extract_score_from_logprobs(self, ret: Dict[str, Any]) -> float:
        """Extract reranking score from generation response with logprobs."""
        # Get logprobs from the response
        meta_info = ret.get("meta_info", {})
        output_top_logprobs = meta_info.get("output_top_logprobs", [])

        # Use output_top_logprobs[0] - the model's prediction for the first generated token
        top_logprobs = output_top_logprobs[0] if output_top_logprobs else []

        # Find yes and no token probabilities
        # Format: list of tuples (logprob, token_id, token_text)
        p_yes = 0.0
        p_no = 0.0
        found_yes = False
        found_no = False

        for item in top_logprobs:
            logprob, token_id = item[0], item[1]
            if token_id == self._yes_token_id:
                p_yes = math.exp(logprob)
                found_yes = True
            elif token_id == self._no_token_id:
                p_no = math.exp(logprob)
                found_no = True
            if found_yes and found_no:
                break

        return _qwen3_rerank_score(p_yes, p_no)

    def _build_rerank_response(
        self, ret: Union[List[Dict[str, Any]], List[float]], request: V1RerankReqInput
    ) -> List[RerankResponse]:
        """Build the rerank response from generation results"""
        responses = []
        for idx, item in enumerate(ret):
            if isinstance(item, dict):
                score_val = item.get("embedding")
                # Some rerank/reward models return scalar score as embedding[0].
                if isinstance(score_val, list):
                    if len(score_val) == 0 or not isinstance(
                        score_val[0], (int, float)
                    ):
                        raise ValueError(
                            f"Invalid embedding score for rerank at index {idx}: {score_val!r}"
                        )
                    score_val = float(score_val[0])
                responses.append(
                    RerankResponse(
                        score=float(score_val),
                        document=(
                            request.documents[idx] if request.return_documents else None
                        ),
                        index=idx,
                        meta_info=item.get("meta_info"),
                    )
                )
            else:
                responses.append(
                    RerankResponse(
                        score=float(item),
                        document=(
                            request.documents[idx] if request.return_documents else None
                        ),
                        index=idx,
                    )
                )

        # When top_n is set, nlargest avoids fully sorting the candidate list
        # (O(N log top_n) vs O(N log N)) — meaningful for large rerank batches.
        # Validator (V1RerankReqInput.validate_top_n) guarantees top_n >= 1.
        if request.top_n is not None:
            return heapq.nlargest(request.top_n, responses, key=lambda x: x.score)

        responses.sort(key=lambda x: x.score, reverse=True)
        return responses