import requests
from bs4 import BeautifulSoup
import os
import json
import gradio as gr
from datasets import Dataset
from PIL import Image
import io
import uuid
import time
import random

DATA_DIR = "/data"
IMAGES_DIR = os.path.join(DATA_DIR, "images")
DATASET_FILE = os.path.join(DATA_DIR, "dataset.json")

# Add a user agent rotation list
USER_AGENTS = [
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.0.3 Safari/605.1.15",
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:89.0) Gecko/20100101 Firefox/89.0"
]

def get_headers(cookies=None):
    headers = {
        "User-Agent": random.choice(USER_AGENTS),
        "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
        "Accept-Language": "en-US,en;q=0.5",
        "Referer": "https://www.google.com/",
        "DNT": "1",
        "Connection": "keep-alive",
        "Upgrade-Insecure-Requests": "1"
    }
    if cookies:
        headers["Cookie"] = cookies
    return headers

def make_request(url, cookies=None):
    time.sleep(random.uniform(1, 3))  # Add a random delay between requests
    return requests.get(url, headers=get_headers(cookies), timeout=10)

def extract_image_url(html_content):
    soup = BeautifulSoup(html_content, 'html.parser')
    
    # First, try to extract the image URL from the <script> tag
    script = soup.find('script', type='text/javascript', string=lambda text: 'image =' in text if text else False)
    
    if script:
        try:
            # Extract and clean the JavaScript object string
            js_object_str = script.string.split('=', 1)[1].strip().rstrip(';')
            
            # Log the string for debugging
            print("Extracted JavaScript object string:", js_object_str)
            
            # Replace single quotes with double quotes to make it a valid JSON string
            js_object_str = js_object_str.replace("'", '"')
            
            # Parse the JSON object
            image_data = json.loads(js_object_str)
            
            # Construct the full image URL
            return f"{image_data['domain']}{image_data['base_dir']}/{image_data['dir']}/{image_data['img']}"
        
        except json.JSONDecodeError as e:
            raise Exception(f"Failed to decode JSON: {str(e)}")
    
    # If the script tag method fails, try to get the image URL from an <img alt> tag
    img_tag = soup.find('img', alt=True)
    if img_tag and 'src' in img_tag.attrs:
        return img_tag['src']
    
    return None

def extract_tags(html_content):
    soup = BeautifulSoup(html_content, 'html.parser')
    
    # Find all list items with the relevant class
    tag_elements = soup.find_all('li', class_='tag-type-general')
    
    tags = []
    for tag_element in tag_elements:
        # The second <a> tag contains the relevant tag name
        tag_links = tag_element.find_all('a')
        if len(tag_links) > 1:
            tag_name = tag_links[1].text
            tags.append(tag_name)
    
    # Join all tags into a single string separated by commas
    return ','.join(tags)

def download_image(url, cookies=None):
    try:
        response = make_request(url, cookies)
        response.raise_for_status()
        return Image.open(io.BytesIO(response.content))
    except requests.RequestException as e:
        raise Exception(f"Failed to download image: {str(e)}")

class DatasetBuilder:
    def __init__(self):
        self.dataset = self.load_dataset()
        os.makedirs(IMAGES_DIR, exist_ok=True)

    def load_dataset(self):
        if os.path.exists(DATASET_FILE):
            with open(DATASET_FILE, 'r') as f:
                return json.load(f)
        return []

    def save_dataset(self):
        with open(DATASET_FILE, 'w') as f:
            json.dump(self.dataset, f)

    def add_image(self, url, cookies=None):
        try:
            response = make_request(url, cookies)
            response.raise_for_status()
            html_content = response.text

            image_url = extract_image_url(html_content)
            if not image_url:
                raise Exception("Failed to extract image URL")

            tags = extract_tags(html_content)
            image = download_image(image_url, cookies)

            # Generate a unique filename
            filename = f"{uuid.uuid4()}.jpg"
            filepath = os.path.join(IMAGES_DIR, filename)

            # Save the image
            image.save(filepath)

            self.dataset.append({
                'image': filename,
                'tags': tags
            })

            self.save_dataset()
            return f"Added image with tags: {tags}"
        except Exception as e:
            return f"Error: {str(e)}"

    def build_huggingface_dataset(self):
        if not self.dataset:
            return "Dataset is empty. Add some images first."

        try:
            hf_dataset = Dataset.from_dict({
                'image': [os.path.join(IMAGES_DIR, item['image']) for item in self.dataset],
                'tags': [item['tags'] for item in self.dataset]
            })
            return "HuggingFace Dataset created successfully!"
        except Exception as e:
            return f"Error creating HuggingFace Dataset: {str(e)}"

    def get_dataset_info(self):
        return f"Current dataset size: {len(self.dataset)} images"

    def get_dataset_preview(self, num_images=5):
        preview = []
        for item in self.dataset[-num_images:]:
            image_path = os.path.join(IMAGES_DIR, item['image'])
            preview.append((image_path, item['tags']))
        return preview

dataset_builder = DatasetBuilder()

def add_image_to_dataset(url, cookies):
    result = dataset_builder.add_image(url, cookies)
    return result, dataset_builder.get_dataset_info(), dataset_builder.get_dataset_preview()

def create_huggingface_dataset():
    return dataset_builder.build_huggingface_dataset()

def view_dataset():
    return dataset_builder.get_dataset_preview(num_images=20)

# Create Gradio interface
with gr.Blocks(theme="huggingface") as iface:
    gr.Markdown("# Image Dataset Builder")
    gr.Markdown("Enter a URL to add an image and its tags to the dataset. Progress is saved automatically.")

    with gr.Row():
        url_input = gr.Textbox(lines=2, placeholder="Enter image URL here...")
        cookies_input = gr.Textbox(lines=2, placeholder="Enter cookies (optional)")
        add_button = gr.Button("Add Image")

    result_output = gr.Textbox(label="Result")
    dataset_info = gr.Textbox(label="Dataset Info", value=dataset_builder.get_dataset_info())

    gr.Markdown("## Dataset Preview")
    preview_gallery = gr.Gallery(label="Recent Additions", show_label=False, elem_id="preview_gallery", columns=5, rows=1, height="auto")

    add_button.click(add_image_to_dataset, inputs=[url_input, cookies_input], outputs=[result_output, dataset_info, preview_gallery])

    create_hf_button = gr.Button("Create HuggingFace Dataset")
    hf_result = gr.Textbox(label="Dataset Creation Result")
    create_hf_button.click(create_huggingface_dataset, inputs=[], outputs=hf_result)

    view_dataset_button = gr.Button("View Dataset")
    dataset_gallery = gr.Gallery(label="Dataset Contents", show_label=False, elem_id="dataset_gallery", columns=5, rows=4, height="auto")
    view_dataset_button.click(view_dataset, inputs=[], outputs=dataset_gallery)

# Launch the interface
iface.launch()