from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch

from app.core.config import REGISTRY


class TranslatorLLM:
    def __init__(self):
        base_path = REGISTRY.get_base_model()
        adapter_path = REGISTRY.get_adapter("llm_t")

        self.tokenizer = AutoTokenizer.from_pretrained(base_path)
        base_model = AutoModelForCausalLM.from_pretrained(
            base_path,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        self.model = PeftModel.from_pretrained(base_model, adapter_path)

    def translate(self, cli_commands: list[str], mode: str) -> list[dict]:
        """
        mode = 'show' -> RESTCONF GET
        mode = 'fix'  -> RESTCONF PATCH/POST
        """
        prompt = self._build_prompt(cli_commands, mode)

        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
        outputs = self.model.generate(
            **inputs,
            max_new_tokens=512,
            do_sample=False
        )

        text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return self._parse_restconf(text)

    def _build_prompt(self, cli_commands: list[str], mode: str) -> str:
        header = (
            "Convert the following Cisco CLI commands into RESTCONF API calls.\n"
            f"Mode: {mode.upper()}\n"
            "Return ONLY valid JSON.\n\n"
        )
        body = "\n".join(cli_commands)
        return header + body

    def _parse_restconf(self, output: str) -> list[dict]:
        """
        Očakáva JSON array RESTCONF callov.
        """
        import json
        try:
            return json.loads(output)
        except Exception:
            raise ValueError(f"Invalid RESTCONF output:\n{output}")