import telebot
import gradio as gr
from gradio_client import Client
import os
import logging
import socks
import socket
import requests
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry
import threading

# Proxy configuration
PROXY_HOST = '64.29.87.81'
PROXY_PORT = 48558
PROXY_USERNAME = 'KNaDX4iRd2t6El0'
PROXY_PASSWORD = 'NleyoQHQJwHbqSH'

# Configure SOCKS proxy
socks.set_default_proxy(socks.SOCKS5, PROXY_HOST, PROXY_PORT, username=PROXY_USERNAME, password=PROXY_PASSWORD)
socket.socket = socks.socksocket

# Create a custom requests session with proxy
session = requests.Session()
session.proxies.update({
    'http': f'socks5://{PROXY_USERNAME}:{PROXY_PASSWORD}@{PROXY_HOST}:{PROXY_PORT}',
    'https': f'socks5://{PROXY_USERNAME}:{PROXY_PASSWORD}@{PROXY_HOST}:{PROXY_PORT}',
})

# Add retry strategy to the session
retry_strategy = Retry(
    total=5,
    backoff_factor=1,
    status_forcelist=[500, 502, 503, 504]
)
adapter = HTTPAdapter(max_retries=retry_strategy)
session.mount('http://', adapter)
session.mount('https://', adapter)

# Initialize the client for image generation
client_image = Client("mukaist/DALLE-4K")

# Define resolutions
resolutions = {
    "896x1152": (896, 1152),
    "1024x1024": (1024, 1024),
    "1216x832": (1216, 832)
}

# Define the default style
DEFAULT_STYLE = "3840 x 2160"

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize the bot with your API token
API_TOKEN = 'YOUR_TELEGRAM_BOT_API_TOKEN'
bot = telebot.TeleBot(API_TOKEN)

def generate_image(prompt, resolution_key, style=DEFAULT_STYLE):
    resolution = resolutions.get(resolution_key, (1024, 1024))
    width, height = resolution
    full_prompt = f"{prompt}, Canon EOS R5, 4K, Photo-Realistic, appearing photorealistic with super fine details, high resolution, natural look, hyper realistic photography, cinematic lighting, --ar 64:37, --v 6.0, --style raw, --stylize 750"
    try:
        result = client_image.predict(
            prompt=full_prompt,
            negative_prompt="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
            use_negative_prompt=True,
            style=style,
            seed=0,
            width=width,
            height=height,
            guidance_scale=5,
            randomize_seed=True,
            api_name="/run"
        )
        logger.info("Image generation successful.")
        return result
    except Exception as e:
        logger.error(f"Error generating image: {e}")
        return None        

def gradio_interface(prompt, resolution_key):
    result = generate_image(prompt, resolution_key)
    
    if result and result[0]:
        file_path = result[0][0].get('image')
        if file_path and os.path.exists(file_path):
            return file_path, "The image was generated successfully."
        else:
            return None, "The image file is not available. Please try again later."
    else:
        return None, "There was an error processing your photo. Please try again later."

def create_gradio_interface(username):
    description = f"Welcome, {username}! Generate images based on prompts."

    with gr.Blocks() as interface:
        gr.HTML(f"<h3>{description}</h3>")
        
        prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
        resolution_dropdown = gr.Dropdown(choices=list(resolutions.keys()), label="Resolution", value="1024x1024")
        generate_button = gr.Button("Generate")
        
        result_output = gr.Image(label="Generated Image", type="pil")
        message_output = gr.Textbox(label="Result", placeholder="Results will be shown here", interactive=False)
        
        generate_button.click(fn=lambda prompt, resolution_key: gradio_interface(prompt, resolution_key),
                              inputs=[prompt_input, resolution_dropdown],
                              outputs=[result_output, message_output])
        
        # Add custom CSS to hide the specific footer element
        gr.HTML("""
        <style>
            footer.svelte-1rjryqp {
                display: none !important;
            }
        </style>
        """)

    return interface

def launch_gradio(username=None):
    if not username:
        username = "Guest"
    # Launch the Gradio interface
    interface = create_gradio_interface(username)
    interface.launch()

@bot.message_handler(commands=['start'])
def handle_start(message):
    # Extract username from message
    username = message.from_user.username or "Guest"
    
    # Construct URL with username
    encoded_url = f"https://measmonysuon-webapp.hf.space/?username={username}"
    
    # Send a message with the inline button to open the web app
    markup = telebot.types.InlineKeyboardMarkup()
    button = telebot.types.InlineKeyboardButton(
        text="Open Web App",
        web_app=telebot.types.WebAppInfo(url=encoded_url)
    )
    markup.add(button)
    
    bot.send_message(
        message.chat.id,
        f"Welcome, {username}! Click the button below to open the web app:",
        reply_markup=markup
    )

def run_telegram_bot():
    bot.polling()

# Run Gradio and Telegram bot concurrently
if __name__ == "__main__":
    # Launch Gradio in a separate thread
    gradio_thread = threading.Thread(target=lambda: launch_gradio("Guest"))
    gradio_thread.start()

    # Start Telegram bot polling
    run_telegram_bot()