from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel
from typing import List, Dict, Any, Union
from base64 import b64decode
from io import BytesIO
import open_clip
import requests
import torch
import numpy as np
from PIL import Image
import uvicorn

app = FastAPI()

class EndpointHandler:
    def __init__(self, path="hf-hub:Styld/marqo-fashionSigLIP"):
        self.model, self.preprocess_train, self.preprocess_val = (
            open_clip.create_model_and_transforms(path)
        )
        
        if torch.cuda.is_available():
            self.model = self.model.cuda()
        
        self.tokenizer = open_clip.get_tokenizer(path)

    def classify_image(self, candidate_labels, image):
        def get_top_prediction(text_probs, labels):
            max_index = text_probs[0].argmax().item()
            return {
                "label": labels[max_index],
                "score": text_probs[0][max_index].item(),
            }

        top_prediction = None
        for i in range(0, len(candidate_labels), 10):
            batch_labels = candidate_labels[i : i + 10]
            image_tensor = self.preprocess_val(image).unsqueeze(0)
            text = self.tokenizer(batch_labels)

            with torch.no_grad(), torch.cuda.amp.autocast():
                image_features = self.model.encode_image(image_tensor)
                text_features = self.model.encode_text(text)
                image_features /= image_features.norm(dim=-1, keepdim=True)
                text_features /= text_features.norm(dim=-1, keepdim=True)

                text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)

            current_top = get_top_prediction(text_probs, batch_labels)
            if top_prediction is None or current_top["score"] > top_prediction["score"]:
                top_prediction = current_top

        return {"label": top_prediction["label"]}

    def combine_embeddings(self, text_embeddings, image_embeddings, text_weight=0.5, image_weight=0.5):
        if text_embeddings is not None:
            avg_text_embedding = np.mean(np.vstack(text_embeddings), axis=0)
        else:
            avg_text_embedding = np.zeros_like(image_embeddings[0])

        if image_embeddings is not None:
            avg_image_embeddings = np.mean(np.vstack(image_embeddings), axis=0)
        else:
            avg_image_embeddings = np.zeros_like(text_embeddings[0])

        combined_embedding = np.average(
            np.vstack((avg_text_embedding, avg_image_embeddings)),
            axis=0,
            weights=[text_weight, image_weight],
        )
        return combined_embedding

    def average_text(self, doc):
        text_chunks = [
            " ".join(doc.split(" ")[i : i + 40])
            for i in range(0, len(doc.split(" ")), 40)
        ]
        text_embeddings = []
        for chunk in text_chunks:
            inputs = self.tokenizer(chunk)
            text_features = self.model.encode_text(inputs)
            text_features /= text_features.norm(dim=-1, keepdim=True)
            text_embeddings.append(text_features.detach().squeeze().numpy())
        combined = self.combine_embeddings(
            text_embeddings, None, text_weight=1, image_weight=0
        )
        return combined

    def embedd_image(self, doc) -> list:
        if not isinstance(doc, str):
            image = doc.get("image")
            if "https://" in image:
                image = image.split("|")
                image = [
                    Image.open(BytesIO(response.content))
                    for response in [requests.get(image) for image in image]
                ][0]
                image = self.preprocess_val(image).unsqueeze(0)
                image_features = self.model.encode_image(image)
                image_features /= image_features.norm(dim=-1, keepdim=True)
                image_embedding = image_features.detach().squeeze().numpy()
                if doc.get("description", "") == "":
                    return image_embedding.tolist()
                else:
                    average_texts = self.average_text(doc.get("description"))
                    combined = self.combine_embeddings(
                        [average_texts],
                        [image_embedding],
                        text_weight=0.5,
                        image_weight=0.5,
                    )
                    return combined.tolist()
        elif isinstance(doc, str):
            return self.average_text(doc).tolist()

    def process_batch(self, batch) -> object:
        try:
            batch = batch.get("batch")
            if not isinstance(batch, list):
                return "Invalid input: batch must be an array of strings.", 400
            embeddings = [self.embedd_image(item) for item in batch]
            return embeddings
        except Exception as e:
            return "An error occurred while processing the request.", 500

    def base64_image_to_pil(self, base64_str) -> Image:
        image_data = b64decode(base64_str)
        image_buffer = BytesIO(image_data)
        image = Image.open(image_buffer)
        return image

handler = EndpointHandler()

class ClassifyRequest(BaseModel):
    candidates: List[str]
    image: str

class EmbeddRequest(BaseModel):
    batch: List[Union[str, Dict[str, str]]]

@app.post("/classify")
def classify(request: ClassifyRequest):
    try:
        image = (
            Image.open(BytesIO(requests.get(request.image).content))
            if "https://" in request.image
            else handler.base64_image_to_pil(request.image)
        )
        response = handler.classify_image(request.candidates, image)
        return response
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/embedd")
def embedd(request: EmbeddRequest):
    try:
        embeddings = handler.process_batch(request.dict())
        return {"embeddings": embeddings}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/process")
async def process(request: Request):
    try:
        data = await request.json()
        if "candidates" in data and "image" in data:
            classify_request = ClassifyRequest(**data)
            return classify(classify_request)
        elif "batch" in data:
            embedd_request = EmbeddRequest(**data)
            return embedd(embedd_request)
        else:
            raise HTTPException(status_code=400, detail="Invalid request format.")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)