from transformers import AutoModelForCausalLM, AutoTokenizer
from jsonschema import validate, ValidationError
from typing import Dict, Any, List
import jsonschema
import torch
import json

class EndpointHandler():
    def __init__(self, path=""):
        self.model = AutoModelForCausalLM.from_pretrained(
            path,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            device_map="auto"
        )
        self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
        self.schema = {
            "type": "object",
            "properties": {
                "objetivoDistancia": {"type": "string"},
                "calentamiento": {"type": "string"},
                "drills": {"type": "array", "items": {"type": "object"}},
                "sprints": {"type": "array", "items": {"type": "object"}},
                "enfriamiento": {"type": "string"}
            },
            "required": ["calentamiento", "drills", "enfriamiento"]
        }

    def parse_and_validate_json(self, output):
        output = output.replace("\n\n", "\n").strip()
        outputs = output.split("\n")
        
        for single_output in outputs:
            try:
                corrected_json = single_output.replace('}\n"', '}, "')
                data = json.loads(corrected_json)
                validate(instance=data, schema=self.schema)
                return True, corrected_json
            except (json.JSONDecodeError, ValidationError):
                continue
        return False, "No se encontró ningún JSON válido según el esquema requerido."

    def model_inference(self, text, **parameters):
        default_parameters = {
            "max_new_tokens": 512,
            "temperature": 0.5,
            "top_p": 0.99,
            "top_k": 20,
            "repetition_penalty": 1.2
        }
        default_parameters.update(parameters)
        inputs = self.tokenizer(text, return_tensors="pt").to("cuda:0")
        inputs_length = len(inputs["input_ids"][0])
        with torch.inference_mode():
            outputs = self.model.generate(**inputs, **default_parameters)
        return self.tokenizer.decode(outputs[0][inputs_length:], skip_special_tokens=True)

    def model_inference_with_validation(self, filled_template, **parameters):
        while True:
            result = self.model_inference(filled_template, **parameters)
            is_valid, response_or_error = self.parse_and_validate_json(result)
            if is_valid:
                return {'generated_text': [response_or_error]}
            else:
                continue

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        input_data = data.pop("inputs", data)
        parameters = data.pop("parameters", {})

        template = """
        [INST] <<SYS>>
        A continuación se presentan parámetros fisiológicos y ambientales capturados por los sensores de un smartwatch utilizado por un atleta. Basándote en estos datos y en tu conocimiento de las rutinas de entrenamiento para atletas, crea una rutina diaria de entrenamiento personalizada para el atleta.
        <</SYS>>
        {data} [/INST]
        """.strip()
        
        data_text = json.dumps(input_data, indent=4)
        filled_template = template.replace("{data}", data_text)

        return self.model_inference_with_validation(filled_template, **parameters)