import endpoint_preflight
from pathlib import Path
from typing import Any

import torch
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast


class EndpointHandler:
    def __init__(self, path: str = ""):
        repo_dir = Path(path)
        for filename in (
            "adapter_config.json", "tokenizer.json", "tokenizer_config.json"
        ):
            if not (repo_dir / filename).is_file():
                raise FileNotFoundError(f"Missing {filename} in {repo_dir}")

        config = PeftConfig.from_pretrained(
            str(repo_dir), local_files_only=True
        )
        if not config.base_model_name_or_path:
            raise ValueError("The adapter configuration does not identify a base model.")

        # Everything in this loading block belongs INSIDE __init__.
        self.tokenizer = PreTrainedTokenizerFast.from_pretrained(
            str(repo_dir), local_files_only=True
        )
        if not self.tokenizer.chat_template:
            raise ValueError("Include the saved GPT-OSS chat template in the repository.")

        # Support both newer and older GPT-OSS special-token spellings.
        vocab = self.tokenizer.get_vocab()
        spellings = (
            ("<|start|>", "<|channel|>", "<|message|>", "<|return|>"),
            ("<|im_start|>", "<|meta_sep|>", "<|im_sep|>", "<|fim_suffix|>"),
        )
        for start, channel, message, finish in spellings:
            if all(token in vocab for token in (start, channel, message, finish)):
                self.final_marker = self.tokenizer.encode(
                    channel + "final" + message, add_special_tokens=False
                )
                self.turn_end_id = vocab[finish]
                self.boundary_ids = {vocab[start], vocab[channel], vocab[message]}
                break
        else:
            raise ValueError("The saved tokenizer is missing GPT-OSS boundary tokens.")

        for token in ("<|end|>", "<|im_end|>", "<|call|>", "<|ghissue|>"):
            if token in vocab:
                self.boundary_ids.add(vocab[token])
        self.boundary_ids.add(self.turn_end_id)

        # Reuse an existing token for padding; do not add vocabulary entries.
        if self.tokenizer.pad_token_id is None:
            if self.tokenizer.eos_token_id is None:
                raise ValueError("The saved tokenizer has neither padding nor EOS configured.")
            self.tokenizer.pad_token = self.tokenizer.eos_token

        base_model = AutoModelForCausalLM.from_pretrained(
            config.base_model_name_or_path,
            revision=getattr(config, "revision", None),
            torch_dtype="auto",
            device_map="auto",
        )
        self.model = PeftModel.from_pretrained(
            base_model,
            str(repo_dir),
            config=config,
            is_trainable=False,
            local_files_only=True,
        )
        self.model.eval()

    def _final_answer(self, prompt_ids: list[int], new_ids: list[int]) -> str:
        # Fail closed: never return an analysis-only or unfinished response.
        if not new_ids or new_ids[-1] != self.turn_end_id:
            raise RuntimeError(
                "Generation did not finish a complete answer. Shorten the request "
                "or increase max_new_tokens within the endpoint's memory budget."
            )

        marker = self.final_marker
        # Include a short prompt tail in case the final header was prefilled.
        ids = prompt_ids[-len(marker):] + new_ids
        positions = [
            i + len(marker)
            for i in range(len(ids) - len(marker) + 1)
            if ids[i:i + len(marker)] == marker
        ]
        if not positions:
            raise RuntimeError("No GPT-OSS final-answer channel was generated.")

        answer_ids = ids[positions[-1]:-1]
        if any(token in self.boundary_ids for token in answer_ids):
            raise RuntimeError("Unexpected message boundaries inside the final answer.")
        answer = self.tokenizer.decode(
            answer_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
        ).strip()
        if not answer:
            raise RuntimeError("The model generated an empty final answer.")
        return answer

    def __call__(self, data: dict[str, Any]) -> list[dict[str, str]]:
        if not isinstance(data, dict):
            raise ValueError("The request body must be a JSON object.")
        prompt = data.get("inputs")
        if not isinstance(prompt, str) or not prompt.strip():
            raise ValueError("'inputs' must be a non-empty prompt string.")

        params = data.get("parameters", {})
        if not isinstance(params, dict):
            raise ValueError("'parameters' must be a JSON object.")
        max_new_tokens = params.get("max_new_tokens", 512)
        if type(max_new_tokens) is not int or not 1 <= max_new_tokens <= 1024:
            raise ValueError("max_new_tokens must be an integer from 1 to 1024.")

        # Your backend can construct this prompt from website selections.
        messages = [{"role": "user", "content": prompt}]
        inputs = self.tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt",
            return_dict=True,
        )
        prompt_ids = inputs["input_ids"][0].tolist()
        context_limit = getattr(self.model.config, "max_position_embeddings", None)
        if context_limit and len(prompt_ids) + max_new_tokens > context_limit:
            raise ValueError("The input plus output budget exceeds the model context limit.")
        inputs = inputs.to(self.model.device)

        with torch.inference_mode():
            outputs = self.model.generate(
                input_ids=inputs["input_ids"],
                attention_mask=inputs["attention_mask"],
                max_new_tokens=max_new_tokens,
                do_sample=False,
                num_beams=1,
                num_return_sequences=1,
                return_dict_in_generate=False,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.turn_end_id,
            )

        new_ids = outputs[0][len(prompt_ids):].tolist()
        return [{"generated_text": self._final_answer(prompt_ids, new_ids)}]