import re
import time
import asyncio
from io import BytesIO
from typing import List, Optional

import httpx
import matplotlib.pyplot as plt
import numpy as np
import torch
import PIL
from transformers import CLIPModel, CLIPProcessor
from PIL import Image


class OffTopicDetector:
    def __init__(self, model_id: str, device: Optional[str] = None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.processor = CLIPProcessor.from_pretrained(model_id)
        self.model = CLIPModel.from_pretrained(model_id).to(self.device)

    def predict_probas(self, images: List[PIL.Image.Image], domain: str,
                valid_templates: Optional[List[str]] = None,
                invalid_classes: Optional[List[str]] = None,
                autocast: bool = True):
        if valid_templates:
            valid_classes = [template.format(domain) for template in valid_templates]
        else:
            valid_classes = [f"a photo of {domain}", f"brochure with {domain} image", f"instructions for {domain}", f"{domain} diagram"]
        if not invalid_classes:
            invalid_classes = ["promotional ad with store information", "promotional text", "google maps screenshot", "business card", "qr code"]
        n_valid = len(valid_classes)
        classes = valid_classes + invalid_classes
        print(f"Valid classes: {valid_classes}", f"Invalid classes: {invalid_classes}", sep="\n")
        n_classes = len(classes)

        start = time.time()
        inputs = self.processor(text=classes, images=images, return_tensors="pt", padding=True).to(self.device)
        if self.device == "cpu" and autocast is True:
            print("Disabling autocast due to device='cpu'.")
            autocast = False
        with torch.autocast(self.device, enabled=autocast):
            with torch.no_grad():
                outputs = self.model(**inputs)
                probas = outputs.logits_per_image.softmax(dim=1).cpu().numpy() # we can take the softmax to get the label probabilities
        end = time.time()
        duration = end - start
        print(f"Device: {self.device}",
              f"Response time: {duration}s",
              f"Response time per image: {round(duration/len(images), 2) * 1000}ms",
              sep="\n")
        valid_probas = probas[:, 0:n_valid].sum(axis=1, keepdims=True)
        invalid_probas = probas[:, n_valid:n_classes].sum(axis=1, keepdims=True)
        return probas, valid_probas, invalid_probas

    def show(self, images: List[PIL.Image.Image], valid_probas: np.ndarray, n_cols: int = 3, title: Optional[str] = None, threshold: Optional[float] = None):
        if threshold is not None:
            prediction = self.apply_threshold(valid_probas, threshold)
            title_scores = [f"Valid: {pred.squeeze()}" for pred in prediction]
        else:
            prediction = np.round(valid_probas[:, 0], 2)
            title_scores = [f"Valid: {pred:.2f}" for pred in prediction]
        n_images = len(images)
        n_rows = int(np.ceil(n_images / n_cols))
        fig, axes = plt.subplots(n_rows, n_cols, figsize=(16, 16))
        for i, ax in enumerate(axes.ravel()):
            ax.axis("off")
            try:
                ax.imshow(images[i])
                ax.set_title(title_scores[i])
            except IndexError:
                continue
        if title:
            fig.suptitle(title)
        fig.tight_layout()
        return

    def predict_item_probas(self, url_or_id: str,
                valid_templates: Optional[List[str]] = None,
                invalid_classes: Optional[List[str]] = None):
        images, domain = self.get_item_data(url_or_id)
        probas, valid_probas, invalid_probas = self.predict_probas(images, domain, valid_templates,
                                                            invalid_classes)
        return images, domain, probas, valid_probas, invalid_probas

    def apply_threshold(self, valid_probas: np.ndarray, threshold: float = 0.4):
        return valid_probas >= threshold

    def get_item_data(self, url_or_id: str):
        if url_or_id.startswith("http"):
            item_id = "".join(url_or_id.split("/")[3].split("-")[:2])
        else:
            item_id = re.sub("-", "", url_or_id)
        response = httpx.get(f"https://api.mercadolibre.com/items/{item_id}").json()
        domain = re.sub("_", " ", response["domain_id"].split("-")[-1]).lower()
        img_urls = [x["url"] for x in response["pictures"]]
        images = self.get_images(img_urls)
        return images, domain

    def get_images(self, urls: List[str]):
        start = time.time()
        images = asyncio.run(self._gather_download_tasks(urls))
        end = time.time()
        duration = end - start
        print(f"Download time: {duration}s",
              f"Download time per image: {round(duration/len(urls), 2) * 1000}ms",
              sep="\n")
        return asyncio.run(self._gather_download_tasks(urls))

    async def _gather_download_tasks(self, urls: List[str]):

        async def _process_download(url: str, client: httpx.AsyncClient):
            response = await client.get(url)
            return Image.open(BytesIO(response.content))

        async with httpx.AsyncClient() as client:
            tasks = [_process_download(url, client) for url in urls]
            return await asyncio.gather(*tasks)

    @staticmethod
    def _non_async_get_item_data(url_or_id: str, save_images: bool = False):
        if url_or_id.startswith("http"):
            item_id = "".join(url_or_id.split("/")[3].split("-")[:2])
        else:
            item_id = re.sub("-", "", url_or_id)
        response = httpx.get(f"https://api.mercadolibre.com/items/{item_id}").json()
        domain = re.sub("_", " ", response["domain_id"].split("-")[-1]).lower()
        img_urls = [x["url"] for x in response["pictures"]]
        images = []
        for img_url in img_urls:
            img = httpx.get(img_url)
            images.append(Image.open(BytesIO(img.content)))
            if save_images:
                with open(re.sub("D_NQ_NP_", "", img_url.split("/")[-1]) , "wb") as f:
                    f.write(img.content)
        return images, domain