@lru_cache(maxsize=1)
def get_tokenizer(self):
    """Load tokenizer if not already loaded."""
    if self.tokenizer is None:
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
    return self.tokenizer

def load_model(self):
    """Load the model with proper error handling and performance optimizations."""
    try:
        if self.model is None:
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                torch_dtype=torch.float16,
                device_map="auto",
                # Add performance optimizations
                low_cpu_mem_usage=True,
                offload_folder="offload",
            )
            
            tokenizer = self.get_tokenizer()
            current_embed_size = self.model.get_input_embeddings().weight.size(0)
            if len(tokenizer) != current_embed_size:
                self.model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=8)
            
        return self.model
    except Exception as e:
        raise RuntimeError(f"Failed to load model: {str(e)}") from e

def get_pipeline(self) -> Pipeline:
    """Get or create the generation pipeline with optimized settings."""
    if self.pipe is None:
        model = self.load_model()
        tokenizer = self.get_tokenizer()
        self.pipe = pipeline(
            task="text-generation",
            model=model,
            tokenizer=tokenizer,
            max_new_tokens=1000,  # Increased for longer responses
            return_full_text=False,
            # Add performance settings
            batch_size=1,
            num_workers=4,
        )
    return self.pipe

def create_prompt(self, user_input: str) -> str:
    """Create a formatted prompt for the model."""
    prompt = dedent(f"""
    Instruction:
    {user_input}
    
    Please provide the response in a valid JSON format with spatial_relationships, task_plan, and performance_metrics.
    """)

    messages = [
        {
            "role": "system",
            "content": "Generate a task plan in JSON format with spatial relationships, task planning, and performance metrics.",
        },
        {"role": "user", "content": prompt},
    ]
    
    return self.get_tokenizer().apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

def parse_json_response(self, text: str) -> Dict[str, Any]:
    """Parse the model output into a JSON object."""
    try:
        # Clean up the text to extract just the JSON part
        # Find the first '{' and last '}'
        start_idx = text.find('{')
        end_idx = text.rfind('}') + 1
        
        if start_idx == -1 or end_idx == 0:
            raise ValueError("No JSON object found in response")
            
        json_str = text[start_idx:end_idx]
        
        # Parse the JSON
        return json.loads(json_str)
    except json.JSONDecodeError as e:
        raise ValueError(f"Failed to parse JSON: {str(e)}\nOriginal text: {text}")

def generate(self, prompt: str) -> Dict[str, Any]:
    """Generate and parse response for a given prompt."""
    try:
        pipe = self.get_pipeline()
        outputs = pipe(prompt)
        raw_response = outputs[0]["generated_text"]
        # return self.parse_json_response(raw_response)
        return raw_response
    except Exception as e:
        raise RuntimeError(f"Generation failed: {str(e)}") from e

try:
    while True:
        user_input = input('Enter prompt (or "quit" to exit): ')
        if user_input.lower() == 'quit':
            break
            
        prompt = manager.create_prompt(user_input)
        print(f'\nFormatted prompt:\n{prompt}\n')
        
        response = manager.generate(prompt)
        
        # Save response to file
        save_response(response)
        
        # Pretty print the response
        print("\nGenerated task plan (formatted):")
        print(json.dumps(response, indent=2))
        
except KeyboardInterrupt:
    print("\nExiting gracefully...")
except Exception as e:
    print(f"An error occurred: {str(e)}")