from typing import Dict, List, Any
from transformers import pipeline
import torch
from datetime import datetime
import pytz
from pathlib import Path

class EndpointHandler:
    def __init__(self, model_path: str = "", tokenizer_path: str = None):
        """
        Initialize the endpoint handler with model and system instruction
        
        Args:
            model_path: Path to the model
            tokenizer_path: Path to the tokenizer (if different from model_path)
        """
        # Set default tokenizer path if not provided
        if tokenizer_path is None:
            tokenizer_path = model_path
            
        # Initialize the pipeline
        self.pipeline = pipeline(
            "text-generation",
            model=model_path,
            tokenizer=tokenizer_path,
            max_length=2000,
            device_map="cuda" if torch.cuda.is_available() else "cpu"
        )
        
        # System instruction
        self.system_instruction =
        """You are an AI that generates a user-friendly user interface based on a user query using the µUI script language. You must output the µUI script covering the sufficient interface that is needed without any other content.

You must include the information mentioned by the user usable to construct your interface.

Only when the request is unclear or require mandatory context or information to generate UI for it, you generate the right UI to collect such missing information or clarify the context.  You also must use µUI (MicroUI) Language for the information collection UI generation as well.

Don't output any other content except µUI script. """
        
        # Default location
        self.default_location = "San Francisco, USA"
        
    def get_formatted_time(self, timezone: str = "America/Los_Angeles") -> str:
        """Get formatted time for given timezone"""
        try:
            tz = pytz.timezone(timezone)
            current_time = datetime.now(tz)
            return current_time.strftime('%Y-%m-%d %H:%M:%S %Z')
        except:
            # Default to PT if timezone is invalid
            tz = pytz.timezone("America/Los_Angeles")
            current_time = datetime.now(tz)
            return current_time.strftime('%Y-%m-%d %H:%M:%S %Z')
    
    def format_message(self, query: str, location: str = None, time: str = None) -> str:
        """Format the user message with location and time"""
        location = location if location else self.default_location
        time = time if time else self.get_formatted_time()
        
        return f"Current User query: {query}\nCurrent Location: {location}\nCurrent Time: {time}"
    
    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Handle the inference request
        
        Args:
            data: Dictionary containing:
                - inputs (str): User query
                - location (str, optional): User location
                - time (str, optional): Custom time
                
        Returns:
            List[Dict]: Model response with metadata
        """
        try:
            # Extract inputs
            if isinstance(data, dict):
                inputs = data.pop("inputs", None)
                if inputs is None:
                    # If no inputs field, treat entire data as input
                    inputs = data
                location = data.pop("location", None)
                time = data.pop("time", None)
            else:
                # If data is not a dict, treat it as the input
                inputs = data
                location = None
                time = None
            
            # Format messages
            messages = [
                {"role": "system", "content": self.system_instruction},
                {"role": "user", "content": self.format_message(inputs, location, time)}
            ]
            
            # Run inference
            output = self.pipeline(messages)
            
            # Extract and format response
            response = {
                "generated_text": output[0]["generated_text"][2]["content"],
                "metadata": {
                    "location": location if location else self.default_location,
                    "time": time if time else self.get_formatted_time()
                }
            }
            
            return [response]
            
        except Exception as e:
            return [{"error": f"Inference failed: {str(e)}"}]