import gradio as gr
import torch
import spaces
import logging
from deep_translator import GoogleTranslator
import pandas as pd
from tqdm import tqdm
import urllib
import aiohttp
import asyncio
from bs4 import BeautifulSoup

# Configure logging to write messages to a file
logging.basicConfig(filename='app.log', level=logging.ERROR)

# Configuration
max_seq_length = 2048
dtype = None  # Auto detection of dtype
load_in_4bit = True  # Use 4-bit quantization to reduce memory usage
peft_model_name = "limitedonly41/website_qwen2_7b_2"

# Initialize model and tokenizer variables
model = None
tokenizer = None

# Async function to fetch data
async def fetch_data(url):
    headers = {
        'Accept': '*/*',
        'Accept-Language': 'ru-RU,ru;q=0.9,en-US;q=0.8,en;q=0.7',
        'Connection': 'keep-alive',
        'Referer': f'{url}',
        'Sec-Fetch-Dest': 'empty',
        'Sec-Fetch-Mode': 'cors',
        'Sec-Fetch-Site': 'cross-site',
        'User-Agent': 'Mozilla/5.0',
        'sec-ch-ua': '"Google Chrome";v="125", "Chromium";v="125", "Not.A/Brand";v="24"',
        'sec-ch-ua-mobile': '?0',
        'sec-ch-ua-platform': '"macOS"',
    }

    encoding = 'utf-8'
    timeout = 10  # Set your desired timeout value in seconds

    try:
        # Function to make the request using urllib
        def get_content():
            req = urllib.request.Request(url, headers=headers)
            with urllib.request.urlopen(req, timeout=timeout) as response:
                return response.read()

        # Async task using executor for blocking I/O
        loop = asyncio.get_event_loop()
        response_content = await loop.run_in_executor(None, get_content)

        soup = BeautifulSoup(response_content, 'html.parser', from_encoding=encoding)
        title = soup.find('title').text if soup.find('title') else ""
        description = soup.find('meta', attrs={'name': 'description'})
        description = description.get("content") if description and "content" in description.attrs else ""

        keywords = soup.find('meta', attrs={'name': 'keywords'})
        keywords = keywords.get("content") if keywords and "content" in keywords.attrs else ""

        h1_all = ". ".join(h.text for h in soup.find_all('h1'))
        paragraphs_all = ". ".join(p.text for p in soup.find_all('p'))
        h2_all = ". ".join(h.text for h in soup.find_all('h2'))
        h3_all = ". ".join(h.text for h in soup.find_all('h3'))

        allthecontent = f"{title} {description} {h1_all} {h2_all} {h3_all} {paragraphs_all}"
        allthecontent = allthecontent[:4999]

        return {
            'url': url,
            'title': title,
            'description': description,
            'keywords': keywords,
            'h1': h1_all,
            'h2': h2_all,
            'h3': h3_all,
            'paragraphs': paragraphs_all,
            'text': allthecontent
        }

    except Exception as e:
        logging.exception(f"Error fetching data for {url}: {e}")
        return {
            'url': url,
            'title': None,
            'description': None,
            'keywords': None,
            'h1': None,
            'h2': None,
            'h3': None,
            'paragraphs': None,
            'text': None
        }

# Main async function to process multiple URLs
async def main(urls):
    tasks = [fetch_data(url) for url in urls]
    results = []
    for future in tqdm(asyncio.as_completed(tasks), total=len(tasks)):
        result = await future
        results.append(result)
    return results

@spaces.GPU()
def classify_website(url):
    global model, tokenizer

    urls = [url]

    # Start asyncio loop for fetching data
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    results_shop = loop.run_until_complete(main(urls))  # Correctly use asyncio loop

    # Convert results to DataFrame
    df_result_train_more = pd.DataFrame(results_shop)
    text = df_result_train_more['text'][0]
    translated = GoogleTranslator(source='auto', target='en').translate(text[:4990])

    try:
        # Load the model and tokenizer if not already loaded
        if model is None or tokenizer is None:
            from unsloth import FastLanguageModel

            model, tokenizer = FastLanguageModel.from_pretrained(
                model_name=peft_model_name,
                max_seq_length=max_seq_length,
                dtype=dtype,
                load_in_4bit=load_in_4bit,
            )
            FastLanguageModel.for_inference(model)

        prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
Categorize the website into one of the 3 categories:

1) OTHER
2) NEWS/BLOG
3) E-commerce

### Input:
{translated}

### Response:"""

        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
        outputs = model.generate(**inputs, max_new_tokens=64, use_cache=True)
        ans = tokenizer.batch_decode(outputs)[0]
        ans_pred = ans.split('### Response:')[1].split('<')[0]

        if 'OTHER' in ans_pred:
            ans_pred = 'OTHER'
        elif 'NEWS/BLOG' in ans_pred:
            ans_pred = 'NEWS/BLOG'
        elif 'E-commerce' in ans_pred:
            ans_pred = 'E-commerce'

        return ans_pred

    except Exception as e:
        logging.exception(e)
        return str(e)

# Create a Gradio interface
iface = gr.Interface(
    fn=classify_website,
    inputs="text",
    outputs="text",
    title="Website Categorization",
    description="Categorize a website into one of the 3 categories: OTHER, NEWS/BLOG, or E-commerce."
)
iface.queue()  # Enable queue with default settings
iface.launch()