import sys
import base64
from io import BytesIO
from PIL import Image
import torch
import numpy as np
from typing import Dict, Any

class EndpointHandler:
    def __init__(self, path: str = ""):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
        # 1. Add the repository directory to Python's sys.path
        # This is strictly required so Python can find 'pipeline.py' and 'upsamplers.py'
        if path not in sys.path:
            sys.path.append(path)
            
        # 2. Import the custom pipeline from the repository
        # Note: Open pipeline.py in your repo to verify the exact class name. 
        # It is usually named something like SegEarthPipeline or OVPipeline.
        from pipeline import SegEarthPipeline 
        
        # 3. Initialize the model. 
        # You can specify the SAR variant (e.g., 'OV-2') or Optical ('OV') here.
        self.model = SegEarthPipeline(model_id=path, variant="OV", device=self.device)

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Expects a JSON payload:
        {
            "inputs": "<base64_encoded_image_string>",
            "prompt": "water body"
        }
        """
        inputs = data.get("inputs")
        prompt = data.get("prompt", "water") # The target class for open-vocabulary grounding

        # Decode base64 image from the FastAPI backend
        if isinstance(inputs, str):
            image_data = base64.b64decode(inputs)
            image = Image.open(BytesIO(image_data)).convert("RGB")
        elif isinstance(inputs, Image.Image):
            image = inputs.convert("RGB")
        else:
            raise ValueError("Invalid image input. Pass a base64 encoded string or PIL Image.")

        # Run inference using their custom pipeline method
        # (Check pipeline.py to see if their method is called .predict() or .__call__())
        with torch.no_grad():
            result_mask = self.model.predict(image, prompt) 
        
        # REST APIs crash if you try to return raw PyTorch Tensors.
        # We must convert the spatial mask to a base64 PNG image string.
        if isinstance(result_mask, torch.Tensor):
            result_mask = result_mask.cpu().numpy()
            
        # Convert binary/normalized mask to a visible 8-bit image
        if isinstance(result_mask, np.ndarray):
            mask_img = Image.fromarray((result_mask * 255).astype(np.uint8))
        else:
            mask_img = result_mask # Fallback if it already returns a PIL Image

        # Encode the mask back to base64
        buffered = BytesIO()
        mask_img.save(buffered, format="PNG")
        mask_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")

        return {
            "status": "success",
            "prompt": prompt,
            "mask_base64": mask_base64
        }