import random
import string
import base64
import time
import requests
import os
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from bs4 import BeautifulSoup
import boto3
from .search_api import SearchAPI


# Utility function to upload the base64 image to S3
def upload_image_to_s3(base64_image, prefix):
    s3 = boto3.client("s3")
    bucket_name = "instant-img-upload"  # S3 버킷 이름을 입력하세요
    base64_image = "data:image/jpeg;base64," + base64_image
    image_data = base64.b64decode(base64_image.split(",")[1])
    # random_number = random.randint(1, 100) + 100 # Google gemma-4b
    random_number = random.randint(1, 100)  # Gemini flash
    object_key = f"{prefix}_query_image_{random_number}.jpg"
    s3.put_object(
        Bucket=bucket_name,
        Key=object_key,
        Body=image_data,
        ContentType="image/jpeg",  # 이미지 MIME 타입 지정
    )
    url = f"https://{bucket_name}.s3.amazonaws.com/{object_key}"
    return url


# Utility function to save base64 image to a file
def save_base64_to_file(base64_image, file_path):
    with open(file_path, "wb") as img_file:
        img_file.write(base64.b64decode(base64_image))


# Utility function to generate a random filename
def generate_random_filename(extension="jpg"):
    random_string = "".join(random.choices(string.ascii_letters + string.digits, k=10))
    return f"{random_string}.{extension}"


# Base SearchAPI class to define search operations
class SearchAPI:
    def search_with_google_lens(self, image_url, max_retries=10, retry_delay=1):
        raise NotImplementedError

    def fetch_and_extract_text(self, url):
        raise NotImplementedError

    def parse_results(self, response_data, question=None, top_k=5):
        raise NotImplementedError


# GoogleLensSearch class extends SearchAPI
class GoogleLensSearch(SearchAPI):
    def __init__(self, model_name="gemma-3-4b-it", api_key="6768f07860d06f6207f5e008"):
        self.model_name = model_name
        self.api_key = api_key
        self.api_url = "https://api.scrapingdog.com/google_lens"

    def encode_image_to_base64(self, img_path) -> str:
        """
        Encode image as base64 string for API input.
        """
        if not os.path.exists(img_path):
            raise FileNotFoundError(f"Image file not found: {img_path}")

        with open(img_path, "rb") as image_file:
            encoded_image = base64.b64encode(image_file.read()).decode("utf-8")

        # Optionally resize the image here if needed (not implemented in this example)
        return encoded_image

    def compute_query(self,base64_image):
        """
        Converts a base64-encoded image into a data URL format.
        """
        image_url = upload_image_to_s3(base64_image, self.model_name)
        return image_url

    def search_with_google_lens(self, image_url, max_retries=10, retry_delay=1):
        """
        Performs a Google Lens search using the ScrapingDog API with retry logic.
        """
        params = {"api_key": self.api_key, "url": image_url}
        for attempt in range(max_retries):
            try:
                response = requests.get(self.api_url, params=params)
                if response.status_code == 200:
                    return response.json()
                else:
                    raise Exception(
                        f"Request failed with status code: {response.status_code}"
                    )
            except Exception as e:
                if attempt < max_retries - 1:
                    print(
                        f"[Retry {attempt + 1}/{max_retries}] Error: {e}. Retrying in {retry_delay} seconds..."
                    )
                    time.sleep(retry_delay)
                else:
                    raise

    def parse_results(self, response_data, question=None, top_k=5):
        """
        Parses the Google Lens results to extract image URLs and summaries.
        """
        parsed_results = []
        if "lens_results" in response_data:
            for result in response_data["lens_results"][:top_k]:
                original_thumbnail = result.get("original_thumbnail", "")
                title = result.get("title", "")
                # link = result.get("link", "")
                # if thumbnail and link:
                #     text_content = self.fetch_and_extract_text(link)
                #     # summary = self.summarize_text_with_openai(text_content)
                parsed_results.append(
                    # {"image_path": original_thumbnail, "caption": title}
                    {
                        "caption": title
                    }  # Only consider text passages as information(since it have errors to process thumbnails
                )
        return parsed_results

    def fetch_and_extract_text(self, url):
        """
        Fetches a webpage using Selenium and extracts visible text.
        """
        try:
            options = Options()
            options.add_argument("--headless")  # Run Chrome in headless mode
            options.add_argument("--disable-gpu")  # Disable GPU acceleration
            driver = webdriver.Chrome(options=options)
            driver.get(url)
            driver.implicitly_wait(10)
            page_source = driver.page_source
            soup = BeautifulSoup(page_source, "html.parser")
            title = soup.find("h1", class_="entry-title").text.strip()
            content = soup.find("div", class_="entry-content").get_text(
                separator=" ", strip=True
            )
            driver.quit()
            return title + content
        except Exception as e:
            headers = {
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
            }
            response = requests.get(url, headers=headers, allow_redirects=True)
            response.raise_for_status()
            soup = BeautifulSoup(response.content, "html.parser")
            text = soup.get_text(separator=" ", strip=True)
            return text

    def retrieve_query(self, image_path, question=None, top_k=5):
        """
        Retrieves the top-k image-text pairs based on a query image.
        """
        # Generate the query image URL
        base64_image = self.encode_image_to_base64(image_path)
        image_data_url = self.compute_query(base64_image)
        image_url = f"https://lens.google.com/uploadbyurl?url={image_data_url}"

        # Perform the search
        response_data = self.search_with_google_lens(image_url)

        # Parse the results
        similar_images = self.parse_results(response_data, question, top_k=top_k)

        # Return the top-k results
        return similar_images[:top_k]
