import gradio as gr
from gradio_client import Client
import os
import logging
import requests
from huggingface_hub import login, whoami, HfFolder, HfApi, SpaceHardware
import time
import torch

webhook_server = os.getenv('webhook_server')
API_TOKEN = os.getenv('ZeroGPU')

#BASE_URL = os.getenv('SpaceURL')
# Get user information
#ser_info = whoami(token=API_TOKEN)

# Set your token
login(token=API_TOKEN, add_to_git_credential=True)
# Initialize API client
api = HfApi(token=API_TOKEN)

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

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

# Get user information to confirm authentication
user_info = whoami(token=API_TOKEN)
print("User Info:", user_info)

#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#print(torch.version.cuda)
#print(f"Is CUDA available: {torch.cuda.is_available()}")
# True
#print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
# Tesla
      
# Define resolutions and default style
resolutions = {
    "896x1152": (896, 1152),
    "1024x1024": (1024, 1024),
    "1216x832": (1216, 832)
}
DEFAULT_STYLE = "3840 x 2160"

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, user_chat_id):
    if not user_chat_id.strip():
        return None, "User Chat ID cannot be blank. Please provide a valid User Chat ID or register with @supportBot to get UID."
    
    # Request large hardware for image generation
    api.request_space_hardware(repo_id='measmonysuon/DALLE-4K', hardware='t4-medium')
    
    # Add a short delay to allow the hardware switch to complete
    time.sleep(5)  # Adjust this value if needed

    gpu_info = {}  # Initialize gpu_info to an empty dictionary
    
    try:
        # Generate the image
        image_info, seed, gpu_info = generate_image(prompt, resolution_key)
        
        if image_info and image_info[0]:
            file_path = image_info[0].get('image')
            if file_path:
                # Combine base URL with file path to create full URL
                base_url = "https://measmonysuon-imagen.hf.space/file="
                full_file_url = base_url + file_path
                
                if os.path.exists(file_path):
                    # Update user points and notify webhook
                    update_user_points(user_chat_id, -5)
                    notify_webhook_with_file(user_chat_id, full_file_url, gpu_info)  # Notify webhook with full file URL
                    
                    return full_file_url, f"Check your telegram, The image was generated successfully. GPU Info: {json.dumps(gpu_info)}"
                
                return None, "The image file is not available. Please try again later."
            
            return None, "The image file is not available. Please try again later."
        
        return None, "There was an error processing your photo. Please try again later."

    finally:
        # Log GPU info before resetting hardware
        logger.info(f"GPU info before reset: {gpu_info}")
        
        # Always request to switch back hardware
        api.request_space_hardware(repo_id='measmonysuon/DALLE-4K', hardware='zero-a10g')
        logger.info("Hardware reset to Zero-A10G")
    
def handle_generate_image(prompt, resolution_key, user_chat_id):
    points = get_user_points(user_chat_id)
    
    if isinstance(points, str):
        if points == "User not found":
            return None, "Please register with https://t.me/myexpsupportBot"
        return None, "Error retrieving points. Please try again later."

    try:
        points = int(points)
    except ValueError:
        return None, "Error processing points. Please try again later."

    if points >= 5:
        result = gradio_interface(prompt, resolution_key, user_chat_id)
        if result[0]:
            # Include the GPU info in the success message
            gpu_info = result[1]  # Assuming gradio_interface returns GPU info in the second part of the tuple
            return result[0], f"Check your telegram, The image was generated successfully. {gpu_info}"
        
        return None, "There was an error processing your photo. Please try again later."
    
    return None, "Insufficient points. Please get more points before generating an image."

def get_gpu_info():
    if torch.cuda.is_available():
        device_name = torch.cuda.get_device_name(0)
        device_properties = torch.cuda.get_device_properties(0)
        total_memory = f"{device_properties.total_memory // (1024 ** 2)} MB"
        return {
            'is_cuda_available': True,
            'device_name': device_name,
            'total_memory': total_memory
        }
    else:
        return {
            'is_cuda_available': False,
            'device_name': 'N/A',
            'total_memory': 'N/A'
        }
        
def request_otp(user_chat_id):
    try:
        response = requests.post(f"{webhook_server}/request_otp", json={"user_chat_id": user_chat_id})
        response.raise_for_status()
        
        logger.debug(f"Response status code: {response.status_code}")
        logger.debug(f"Response content: {response.text}")
        
        response_data = response.json()
        
        if 'error' in response_data:
            logger.error(f"Error response from /request_otp: {response_data['error']}")
            return {"success": False, "error": response_data.get("error", "Unknown error")}
        
        return {"success": response_data.get("success", False)}
    
    except requests.RequestException as e:
        logger.error(f"Request exception: {e}")
        return {"success": False, "error": "Request exception occurred"}
    
    except ValueError as e:
        logger.error(f"JSON decode error: {e}")
        return {"success": False, "error": "Invalid JSON response"}

def verify_otp(user_chat_id, otp):
    try:
        response = requests.post(f"{webhook_server}/verify_otp", json={"user_chat_id": user_chat_id, "otp": otp})
        response.raise_for_status()
        return response.json().get("valid", False)
    except requests.RequestException as e:
        logger.error(f"Request exception: {e}")
        return False
    except ValueError as e:
        logger.error(f"JSON decode error: {e}")
        return False

def retrieve_user_points(user_chat_id, otp):
    if not verify_otp(user_chat_id, otp):
        return None, "Invalid or expired OTP. Please request a new OTP."

    response = requests.get(f"{webhook_server}/get_points", params={"user_chat_id": user_chat_id})
    if response.status_code == 200:
        data = response.json()
        points = data.get("points", 0)
        return points, "OTP verification and points retrieved successfully. Now, enter pormpt and generate image"
    return None, "Error retrieving points. Please try again later."

def update_user_points(user_chat_id, points):
    webhook_url = f"{webhook_server}/update_user_points"
    payload = {'user_chat_id': user_chat_id, 'points': points}
    try:
        response = requests.post(webhook_url, json=payload)
        response.raise_for_status()
        logger.info(f"User with ID {user_chat_id} had their points updated by {points}. Webhook response: {response.status_code}")
    except requests.RequestException as e:
        logger.error(f"Error updating user points via webhook: {e}")

def get_user_points(user_chat_id):
    try:
        response = requests.get(f"{webhook_server}/get_points", params={'user_chat_id': user_chat_id})
        response.raise_for_status()
        data = response.json()
        if 'error' in data:
            return "User not found"
        return data.get('points', 0)
    except requests.RequestException as e:
        logger.error(f"Error retrieving user points: {e}")
        return "Error retrieving points"

def notify_webhook(user_chat_id):
    webhook_url = f"{webhook_server}/notify_update"
    payload = {'user_chat_id': user_chat_id}
    try:
        response = requests.post(webhook_url, json=payload)
        response.raise_for_status()
        logger.info(f"Webhook notification sent successfully. Response: {response.status_code}")
    except requests.RequestException as e:
        logger.error(f"Error sending webhook notification: {e}")

def notify_webhook_with_file(user_chat_id, full_file_url, gpu_info):
    webhook_url = f"{webhook_server}/notify_image"
    payload = {
        'user_chat_id': user_chat_id,
        'file_path': full_file_url,
        'gpu_info': gpu_info
    }
    try:
        response = requests.post(webhook_url, json=payload)
        response.raise_for_status()
        logger.info(f"Webhook notification sent successfully for user {user_chat_id}. Response: {response.status_code}")
    except requests.RequestException as e:
        logger.error(f"Error sending webhook notification with file URL: {e}")
'''        
def gradio_interface(prompt, resolution_key, user_chat_id):
    if not user_chat_id.strip():
        return None, "User Chat ID cannot be blank. Please provide a valid User Chat ID."

    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):
            update_user_points(user_chat_id, -5)
            notify_webhook_with_file(user_chat_id, file_path)  # Notify webhook with file path
            return file_path, "The image was generated successfully."
        return None, "The image file is not available. Please try again later."
    return None, "There was an error processing your photo. Please try again later."
'''

def request_otp_and_update_ui(user_chat_id):
    # Validate user_chat_id
    if not user_chat_id.strip():
        return "User Chat ID cannot be blank. Please provide a valid User Chat ID or register with @supportBot to get UID."

    otp_response = request_otp(user_chat_id)
    if otp_response.get("success"):
        return "OTP has been sent to your Telegram"
    return "Error requesting OTP. Please try again later."

# Define HTML content directly
PRIVACY_POLICY_HTML = """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Privacy Policy</title>
</head>
<body>
    <h1>Privacy Policy</h1>
    <p><strong>Effective Date:</strong>15 August, 2024</p>

    <h2>1. Introduction</h2>
      <p>Welcome to <a href="https://t.me/myexpsupportBot">@supportBot</a> and Its Mini App. We value your privacy and are committed to protecting your personal data. This Privacy Policy explains how we collect, use, and secure your information when you interact with our Telegram bot and mini app.</p>

    <h2>2. Information We Collect</h2>
    <p>We collect the following types of information:</p>
    <ul>
        <li><strong>Telegram Bot:</strong>
            <ul>
                <li><strong>User ID:</strong> A unique identifier for users of the Telegram bot.</li>
                <li><strong>Chat History:</strong> Messages exchanged between the user and the bot.</li>
                <li><strong>Usage Data:</strong> Information about how the bot is used, including commands and interactions.</li>
            </ul>
        </li>
        <li><strong>Mini App:</strong>
            <ul>
                <li><strong>User Data:</strong> Information provided directly by users, such as names, email addresses, or other personal details.</li>
                <li><strong>App Usage Data:</strong> Details on how the mini app is used, including features accessed and interactions.</li>
            </ul>
        </li>
    </ul>

    <h2>3. How We Use Your Information</h2>
    <p>We use your information for the following purposes:</p>
    <ul>
        <li><strong>To Provide and Improve Our Services:</strong> To enhance the functionality and performance of our Telegram bot and mini app.</li>
        <li><strong>To Communicate with You:</strong> To send updates, respond to inquiries, and provide support.</li>
        <li><strong>To Analyze Usage:</strong> To understand how our services are used and to identify areas for improvement.</li>
    </ul>

    <h2>4. Data Security</h2>
    <p>We implement appropriate technical and organizational measures to protect your data from unauthorized access, alteration, disclosure, or destruction. These measures include:</p>
    <ul>
        <li><strong>Encryption:</strong> Data is encrypted during transmission to safeguard your information.</li>
        <li><strong>Access Controls:</strong> Only authorized personnel have access to your data.</li>
        <li><strong>Regular Audits:</strong> We conduct regular security audits to ensure our measures are effective.</li>
    </ul>

    <h2>5. Data Retention</h2>
    <p>We retain your data only as long as necessary to fulfill the purposes outlined in this Privacy Policy or as required by law. Once the data is no longer needed, it will be securely deleted or anonymized.</p>

    <h2>6. Third-Party Services</h2>
    <p>We may use third-party services to assist in providing our Telegram bot and mini app. These third parties are contractually obligated to protect your data and use it only for the purposes we specify.</p>

    <h2>7. Your Rights</h2>
    <p>You have the right to:</p>
    <ul>
        <li><strong>Access:</strong> Request a copy of the personal data we hold about you.</li>
        <li><strong>Rectify:</strong> Request correction of any inaccurate or incomplete data.</li>
        <li><strong>Delete:</strong> Request deletion of your data under certain conditions.</li>
        <li><strong>Object:</strong> Object to the processing of your data for specific purposes.</li>
    </ul>
    <p>To exercise these rights, please contact us at <strong>Email:</strong>measmonysuon@gmail.com.</p>

    <h2>8. Changes to This Privacy Policy</h2>
    <p>We may update this Privacy Policy from time to time. Any changes will be posted on this page, and we will notify you of significant updates. Your continued use of our services after changes indicates your acceptance of the updated policy.</p>

    <h2>9. Contact Us</h2>
    <p>If you have any questions or concerns about this Privacy Policy or our data practices, please contact us at:</p>
    <ul>
        <li><strong>Email:</strong>measmonysuon@gmail.com</li>
        <li><strong>Telegram:</strong>@measmony</li>
    </ul>
    <p>Thank you for using <a href="https://t.me/myexpsupportBot">@supportBot</a> and Its Mini App. We are committed to protecting your privacy and providing a safe and secure experience.</p>
</body>
</html>
"""

def create_gradio_interface():
    with gr.Blocks() as interface:
        gr.HTML("""
        <h1>AI ImaGen DALL·E 4K</h1>
        <h3>Welcome! Please enter your UID and OTP to continue or Register with <a href="https://t.me/myexpsupportBot">@SupportBot</a> to get UID</h3>
        <p>Read more about our PRIVACY POLICY below</a></p>
        """)

        with gr.Row():
            user_chat_id_input = gr.Textbox(label="Your UID", placeholder="Enter your UID")
            user_otp_input = gr.Textbox(label="OTP", type="password", placeholder="Enter your OTP")
            points_output = gr.Textbox(label="Your Balance", placeholder="Your Points", interactive=False)
            
        with gr.Row():
            request_otp_button = gr.Button("Request OTP")
            get_points_button = gr.Button("Verify OTP")

        message_output = gr.Textbox(label="Message", placeholder="Results and notifications will be shown here", interactive=False)
        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")
        
        generate_button.click(
            fn=handle_generate_image,
            inputs=[prompt_input, resolution_dropdown, user_chat_id_input],
            outputs=[result_output, message_output]
        )

        get_points_button.click(
            fn=retrieve_user_points,
            inputs=[user_chat_id_input, user_otp_input],
            outputs=[points_output, message_output]
        )
        
        request_otp_button.click(
            fn=request_otp_and_update_ui,
            inputs=[user_chat_id_input],
            outputs=[message_output]
        )

         # Display Privacy Policy
        with gr.Row():
            gr.HTML("""
            <button onclick="document.getElementById('privacy-policy').style.display='block'">View PRIVACY POLICY</button>
            <div id="privacy-policy" style="display:none;">
                {}
            </div>
            """.format(PRIVACY_POLICY_HTML))
            
        gr.HTML("""
        <style>
            footer.svelte-1rjryqp {
                display: none !important;
            }
        </style>
        """)

    return interface

def launch_gradio():
    interface = create_gradio_interface()
    interface.queue(max_size=20, api_open=False)  # Configure queue and API accessibility
    interface.launch(show_error=True)  # Enable verbose error reporting

if __name__ == '__main__':
    launch_gradio()