# Copyright 2023-2024 SGLang Team
# 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.
# ==============================================================================

from __future__ import annotations

import logging
import time
import uuid
from typing import TYPE_CHECKING, Dict, Optional

from sglang.srt.managers.io_struct import (
    CloseSessionReqInput,
    OpenSessionReqInput,
    OpenSessionReqOutput,
    TokenizedGenerateReqInput,
)
from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req
from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache

if TYPE_CHECKING:
    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache

logger = logging.getLogger(__name__)


class SessionReqNode:
    def __init__(
        self,
        req: Req,
        parent: Optional["SessionReqNode"] = None,
        children=None,
    ):
        self.req = req
        self.parent = parent
        if parent is not None:
            parent.children.append(self)
        self.children = [] if not children else children

    def clear_children(self, req_dict):
        for req_node in self.children:
            req_node.clear(req_dict)
        self.children = []

    def clear(self, req_dict):
        for req_node in self.children:
            req_node.clear(req_dict)

        if self.req.finished_reason is None:
            self.req.to_finish = FINISH_ABORT()
        del req_dict[self.req.rid]

    def abort(self):
        if self.req.finished_reason is None:
            self.req.to_finish = FINISH_ABORT()

    def __str__(self):
        return self._str_helper(self.req.rid)

    def _str_helper(self, prefix=""):
        if len(self.children) == 0:
            return prefix + "\n"
        else:
            origin_prefix = prefix
            prefix += " -- " + self.children[0].req.rid
            ret = self.children[0]._str_helper(prefix)
            for child in self.children[1:]:
                prefix = " " * len(origin_prefix) + " \\- " + child.req.rid
                ret += child._str_helper(prefix)
            return ret


class Session:
    def __init__(
        self,
        capacity_of_str_len: int,
        session_id: Optional[str] = None,
        streaming: bool = False,
        timeout: Optional[float] = None,
    ):
        self.session_id = session_id if session_id is not None else uuid.uuid4().hex
        self.capacity_of_str_len = capacity_of_str_len
        self.streaming = streaming
        self.timeout = timeout
        self.last_active_time: float = time.monotonic()
        self.req_nodes: Dict[str, SessionReqNode] = {}

    def is_timed_out(self) -> bool:
        if self.timeout is None:
            return False
        return time.monotonic() - self.last_active_time > self.timeout

    def create_req(
        self,
        req: TokenizedGenerateReqInput,
        tokenizer,
        vocab_size: int,
        eos_token_ids=None,
    ):
        assert req.session_params is not None
        self.last_active_time = time.monotonic()
        session_params = req.session_params

        last_req_node = None
        last_req = None
        abort = False
        abort_message = ""
        if self.streaming:
            # Streaming sessions: only simple appends allowed; reject otherwise.
            if session_params.replace:
                abort = True
                abort_message = "Streaming sessions do not support replace."
            elif session_params.drop_previous_output:
                abort = True
                abort_message = (
                    "Streaming sessions do not support drop_previous_output."
                )
            elif session_params.offset and session_params.offset != 0:
                abort = True
                abort_message = "Streaming sessions do not support offset."
            elif self.req_nodes:
                assert len(self.req_nodes) == 1
                _, last_req_node = self.req_nodes.popitem()
                last_req = last_req_node.req
        elif session_params.replace:
            if session_params.rid is None:
                for _, req_node in self.req_nodes.items():
                    req_node.clear(self.req_nodes)
            else:
                if session_params.rid not in self.req_nodes:
                    abort = True
                    abort_message = "Invalid request session id"
                else:
                    last_req_node = self.req_nodes[session_params.rid]
                    last_req_node.abort()
                    last_req = last_req_node.req
                    last_req_node.clear_children(self.req_nodes)
        else:
            if session_params.rid is not None:
                if session_params.rid not in self.req_nodes:
                    abort = True
                    abort_message = "Invalid request session id"
                else:
                    last_req_node = self.req_nodes[session_params.rid]
                    last_req = last_req_node.req
                    if not last_req.finished():
                        abort = True
                        abort_message = "Session request is appending to a request that hasn't finished."
                        logging.warning(abort_message)

        if last_req is not None:
            # trim bos token if it is an append
            if (
                tokenizer is not None
                and req.input_ids
                and req.input_ids[0] == tokenizer.bos_token_id
            ):
                req.input_ids = req.input_ids[1:]
                # Adjust mm_item offsets since they were computed on
                # the pre-strip sequence (with BOS at position 0)
                if req.mm_inputs:
                    for item in req.mm_inputs.mm_items:
                        if item.offsets:
                            if any(s == 0 for s, _ in item.offsets):
                                logging.warning(
                                    "mm_item offset starts at 0 (BOS position), "
                                    "clamping to 0 after BOS strip"
                                )
                            item.offsets = [
                                (max(0, s - 1), max(0, e - 1)) for s, e in item.offsets
                            ]

            input_ids = (
                last_req.origin_input_ids
                + last_req.output_ids[: last_req.sampling_params.max_new_tokens]
            )

            if session_params.drop_previous_output:
                input_ids = last_req.origin_input_ids[:]

            if session_params.offset and session_params.offset != 0:
                input_ids = input_ids[: session_params.offset] + req.input_ids
            else:
                input_ids += req.input_ids

            input_ids_unpadded = (
                last_req.origin_input_ids_unpadded
                + last_req.output_ids[: last_req.sampling_params.max_new_tokens]
            )
            if session_params.drop_previous_output:
                input_ids_unpadded = last_req.origin_input_ids_unpadded[:]

            if session_params.offset and session_params.offset != 0:
                input_ids_unpadded = (
                    input_ids_unpadded[: session_params.offset] + req.input_ids
                )
            else:
                input_ids_unpadded += req.input_ids
        else:
            input_ids = req.input_ids
            input_ids_unpadded = req.input_ids

        new_req = Req(
            rid=req.rid,
            origin_input_text=None,
            origin_input_ids=input_ids,
            origin_input_ids_unpadded=input_ids_unpadded,
            sampling_params=req.sampling_params,
            lora_id=req.lora_id,
            session=self,
            custom_logit_processor=req.custom_logit_processor,
            stream=req.stream,
            return_logprob=req.return_logprob,
            top_logprobs_num=req.top_logprobs_num,
            token_ids_logprob=req.token_ids_logprob,
            vocab_size=vocab_size,
            eos_token_ids=eos_token_ids,
            require_reasoning=req.require_reasoning,
            return_hidden_states=req.return_hidden_states,
            return_routed_experts=req.return_routed_experts,
            priority=req.priority,
            routing_key=req.routing_key,
            http_worker_ipc=req.http_worker_ipc,
            bootstrap_host=req.bootstrap_host,
            bootstrap_port=req.bootstrap_port,
            bootstrap_room=req.bootstrap_room,
            routed_dp_rank=req.routed_dp_rank,
            disagg_prefill_dp_rank=req.disagg_prefill_dp_rank,
            extra_key=req.extra_key,
            time_stats=req.time_stats,
        )
        if last_req is not None:
            new_req.multimodal_inputs = last_req.multimodal_inputs
        new_req.tokenizer = tokenizer

        if abort:
            new_req.set_finish_with_abort(abort_message)
        elif self.streaming:
            if last_req is not None:
                last_req.session = None
            self.req_nodes[req.rid] = SessionReqNode(new_req)
        else:
            new_req_node = SessionReqNode(new_req, last_req_node)
            self.req_nodes[req.rid] = new_req_node

        return new_req


class SessionController:
    def __init__(self, tree_cache: BasePrefixCache):
        self.sessions: Dict[str, Session] = {}
        self._last_reap_time: float = 0.0
        self.tree_cache = tree_cache

    def __contains__(self, session_id: str) -> bool:
        return session_id in self.sessions

    def get(self, session_id: str) -> Optional[Session]:
        return self.sessions.get(session_id)

    def open(self, recv_req: OpenSessionReqInput) -> OpenSessionReqOutput:
        session_id = recv_req.session_id
        if session_id in self.sessions:
            logger.warning(f"session id {session_id} already exist, cannot open.")
            return OpenSessionReqOutput(session_id, False)
        elif session_id is None:
            logger.warning("session id is None, cannot open.")
            return OpenSessionReqOutput(session_id, False)
        else:
            self.sessions[session_id] = Session(
                recv_req.capacity_of_str_len,
                session_id,
                streaming=bool(recv_req.streaming),
                timeout=recv_req.timeout,
            )
            return OpenSessionReqOutput(session_id, True)

    def close(self, recv_req: CloseSessionReqInput):
        session_id = recv_req.session_id
        if session_id not in self.sessions:
            logger.warning(f"session id {session_id} does not exist, cannot delete.")
        else:
            self._close(session_id)

    def _close(self, session_id: str) -> int:
        session = self.sessions[session_id]
        if session.streaming and session.req_nodes:
            assert len(session.req_nodes) == 1
            req = next(iter(session.req_nodes.values())).req
            if not req.finished():
                req.session = None

        # Release multimodal features held by session requests.
        # Session reqs skip the normal mm cleanup path (scheduler and
        # output_processor) so features stay alive until the session closes.
        seen_mm = set()
        for node in session.req_nodes.values():
            mm = node.req.multimodal_inputs
            if mm is not None and id(mm) not in seen_mm:
                seen_mm.add(id(mm))
                mm.release_features()
            node.req.multimodal_inputs = None

        freed_tokens = 0
        if isinstance(self.tree_cache, SessionAwareCache):
            freed_tokens = self.tree_cache.release_session(session_id)
        del self.sessions[session_id]
        return freed_tokens

    def _is_idle_streaming_session(self, session: Session) -> bool:
        if not session.streaming:
            return False
        if not session.req_nodes:
            return False
        return all(node.req.finished() for node in session.req_nodes.values())

    def get_streaming_session_cache_status(
        self, session_id: Optional[str] = None
    ) -> Dict[str, object]:
        if not isinstance(self.tree_cache, SessionAwareCache):
            return {
                "enabled": False,
                "eviction_policy": None,
                "session_count": 0,
                "resident_session_count": 0,
                "held_tokens": 0,
                "available_tokens": 0,
                "capacity_tokens": 0,
                "idle_evictable_session_count": 0,
                "idle_evictable_tokens": 0,
                "sessions": [],
                "target_session": None,
            }

        all_statuses = self.tree_cache.list_session_statuses()
        session_map = {status["session_id"]: status for status in all_statuses}
        idle_evictable_tokens = 0
        idle_evictable_session_count = 0
        sessions = []
        for status in all_statuses:
            session = self.sessions.get(status["session_id"])
            idle_evictable = session is not None and self._is_idle_streaming_session(
                session
            )
            if idle_evictable and status["resident"]:
                idle_evictable_session_count += 1
                idle_evictable_tokens += int(status["resident_tokens"])
            if session_id is None or status["session_id"] == session_id:
                sessions.append(
                    {
                        **status,
                        "streaming": bool(session.streaming) if session is not None else True,
                        "idle_evictable": idle_evictable,
                        "timed_out": session.is_timed_out() if session is not None else False,
                    }
                )

        target_session = session_map.get(session_id) if session_id is not None else None
        return {
            "enabled": True,
            "eviction_policy": "lru",
            "session_count": len(all_statuses),
            "resident_session_count": sum(1 for status in all_statuses if status["resident"]),
            "held_tokens": int(self.tree_cache.session_held_tokens()),
            "available_tokens": int(self.tree_cache.token_to_kv_pool_allocator.available_size()),
            "capacity_tokens": int(self.tree_cache.token_to_kv_pool_allocator.size),
            "idle_evictable_session_count": idle_evictable_session_count,
            "idle_evictable_tokens": int(idle_evictable_tokens),
            "sessions": sessions,
            "target_session": (
                {
                    **target_session,
                    "idle_evictable": (
                        self._is_idle_streaming_session(self.sessions[target_session["session_id"]])
                        if target_session is not None
                        and target_session["session_id"] in self.sessions
                        else False
                    ),
                }
                if target_session is not None
                else None
            ),
        }

    def evict_idle_streaming_sessions(
        self,
        required_tokens: int = 0,
        max_sessions: Optional[int] = None,
    ) -> Dict[str, int]:
        return self.evict_idle_streaming_sessions_lru(
            required_tokens=required_tokens,
            max_sessions=max_sessions,
        )

    def evict_idle_streaming_sessions_lru(
        self,
        required_tokens: int = 0,
        min_available_tokens: int = 0,
        target_available_tokens: Optional[int] = None,
        max_sessions: Optional[int] = None,
        exclude_session_ids: Optional[set[str]] = None,
    ) -> Dict[str, int]:
        if not isinstance(self.tree_cache, SessionAwareCache):
            return {
                "evicted_session_count": 0,
                "freed_tokens": 0,
                "available_tokens_before": 0,
                "available_tokens_after": 0,
            }

        available_tokens_before = int(
            self.tree_cache.token_to_kv_pool_allocator.available_size()
        )
        evicted_session_count = 0
        freed_tokens = 0
        target_available_tokens = max(
            int(min_available_tokens),
            int(target_available_tokens or 0),
        )
        candidates = []
        for status in self.tree_cache.list_session_statuses():
            if exclude_session_ids and status["session_id"] in exclude_session_ids:
                continue
            session = self.sessions.get(status["session_id"])
            if session is None or not self._is_idle_streaming_session(session):
                continue
            if not status["resident"]:
                continue
            candidates.append(status)

        for status in candidates:
            if max_sessions is not None and evicted_session_count >= max_sessions:
                break
            available_tokens_now = available_tokens_before + freed_tokens
            if (
                required_tokens > 0
                and freed_tokens >= required_tokens
                and available_tokens_now >= min_available_tokens
            ):
                break
            if (
                required_tokens <= 0
                and min_available_tokens <= 0
                and target_available_tokens <= 0
            ):
                break
            if (
                required_tokens <= 0
                and available_tokens_now >= min_available_tokens
                and available_tokens_now >= target_available_tokens
            ):
                break

            freed = self._close(status["session_id"])
            evicted_session_count += 1
            freed_tokens += int(freed)

        return {
            "evicted_session_count": evicted_session_count,
            "freed_tokens": freed_tokens,
            "available_tokens_before": available_tokens_before,
            "available_tokens_after": available_tokens_before + freed_tokens,
        }

    def maybe_reap(self, now: float, interval: float = 1.0):
        # reap sessions every second
        if now - self._last_reap_time > interval:
            self._last_reap_time = now
            timed_out = [
                sid for sid, session in self.sessions.items() if session.is_timed_out()
            ]
            for sid in timed_out:
                logger.info(f"Session {sid} timed out, closing.")
                self._close(sid)

    @staticmethod
    def adjust_mm_offsets(recv_req: TokenizedGenerateReqInput, req: Req, image_inputs):
        # For session requests, adjust mm_inputs offsets by the prefix length.
        # Session.create_req prepends previous context to origin_input_ids,
        # so offsets from the new prompt need to be shifted.
        if len(recv_req.input_ids) >= len(req.origin_input_ids):
            return
        prefix_len = len(req.origin_input_ids) - len(recv_req.input_ids)
        for mm_item in image_inputs.mm_items:
            if mm_item.offsets:
                mm_item.offsets = [
                    (start + prefix_len, end + prefix_len)
                    for start, end in mm_item.offsets
                ]
