import gradio as gr
from backend.lcm_text_to_image import LCMTextToImage
from backend.models.lcmdiffusion_setting import LCMLora, LCMDiffusionSetting
from constants import DEVICE, LCM_DEFAULT_MODEL_OPENVINO
from time import perf_counter
import numpy as np
from cv2 import imencode
import base64
from backend.device import get_device_name
from constants import APP_VERSION
from backend.device import is_openvino_device
from PIL import Image
from backend.models.lcmdiffusion_setting import DiffusionTask
from backend.safety_check import is_safe_image
from pprint import pprint
from transformers import pipeline
import os

lcm_text_to_image = LCMTextToImage()
lcm_lora = LCMLora(
    base_model_id="Lykon/dreamshaper-7",
    lcm_lora_id="latent-consistency/lcm-lora-sdv1-5",
)
classifier = pipeline(
    "image-classification",
    model="Falconsai/nsfw_image_detection",
)


# https://github.com/gradio-app/gradio/issues/2635#issuecomment-1423531319
def encode_pil_to_base64_new(pil_image):
    image_arr = np.asarray(pil_image)[:, :, ::-1]
    _, byte_data = imencode(".png", image_arr)
    base64_data = base64.b64encode(byte_data)
    base64_string_opencv = base64_data.decode("utf-8")
    return "data:image/png;base64," + base64_string_opencv


# monkey patching encode pil
gr.processing_utils.encode_pil_to_base64 = encode_pil_to_base64_new


def get_logo_html():
    """Get the AEGIS logo HTML with proper path handling"""
    logo_path = "logo/aegis.jpeg"
    
    # Check if logo exists
    if os.path.exists(logo_path):
        # For local development, use relative path
        return f"""
        <div style="text-align: center; margin-bottom: 20px;">
            <img src="{logo_path}" alt="AEGIS Logo" style="max-width: 200px; height: auto; border-radius: 10px; box-shadow: 0 4px 15px rgba(30, 60, 114, 0.3); transition: all 0.3s ease;" onmouseover="this.style.transform='scale(1.05)'; this.style.boxShadow='0 6px 20px rgba(30, 60, 114, 0.5)';" onmouseout="this.style.transform='scale(1)'; this.style.boxShadow='0 4px 15px rgba(30, 60, 114, 0.3)';">
        </div>
        """
    else:
        # Fallback with emoji if logo not found
        return """
        <div style="text-align: center; margin-bottom: 20px;">
            <div style="font-size: 80px; margin-bottom: 10px;">🧬</div>
            <p style="color: #1e3c72; font-weight: bold; margin: 0;">AEGIS Bio-Digital Lab 10</p>
        </div>
        """


def predict(
    prompt,
    steps,
    seed,
    use_seed,
):
    print(f"prompt - {prompt}")
    lcm_diffusion_setting = LCMDiffusionSetting()
    lcm_diffusion_setting.lcm_model_id = "rupeshs/hyper-sd-sdxl-1-step"
    lcm_diffusion_setting.diffusion_task = DiffusionTask.text_to_image.value
    lcm_diffusion_setting.openvino_lcm_model_id = "rupeshs/sd-turbo-openvino"
    lcm_diffusion_setting.use_lcm_lora = False
    lcm_diffusion_setting.prompt = prompt
    lcm_diffusion_setting.guidance_scale = 1.0
    lcm_diffusion_setting.inference_steps = steps
    lcm_diffusion_setting.seed = seed
    lcm_diffusion_setting.use_seed = use_seed
    lcm_diffusion_setting.use_safety_checker = True
    lcm_diffusion_setting.use_tiny_auto_encoder = False
    # lcm_diffusion_setting.image_width = 320 if is_openvino_device() else 512
    # lcm_diffusion_setting.image_height = 320 if is_openvino_device() else 512
    lcm_diffusion_setting.image_width = 512
    lcm_diffusion_setting.image_height = 512
    lcm_diffusion_setting.use_openvino = True
    lcm_diffusion_setting.use_tiny_auto_encoder = True
    pprint(lcm_diffusion_setting.model_dump())
    lcm_text_to_image.init(lcm_diffusion_setting=lcm_diffusion_setting)
    start = perf_counter()
    images = lcm_text_to_image.generate(lcm_diffusion_setting)
    latency = perf_counter() - start
    print(f"Latency: {latency:.2f} seconds")
    result = images[0]
    if is_safe_image(
        classifier,
        result,
    ):
        return result  # .resize([512, 512], Image.LANCZOS)
    else:
        print("Unsafe image detected")
        return Image.new("RGB", (512, 512), (0, 0, 0))


css = """
#container{
    margin: 0 auto;
    max-width: 55rem;
    background: rgba(255, 255, 255, 0.95);
    border-radius: 15px;
    padding: 25px;
    box-shadow: 0 8px 32px rgba(0,0,0,0.1);
    backdrop-filter: blur(10px);
    border: 2px solid rgba(30, 60, 114, 0.3);
}
#intro{
    max-width: 100%;
    text-align: center;
    margin: 0 auto;
    background: linear-gradient(135deg, #1e3c72 0%, #2a5298 100%);
    color: white;
    padding: 25px;
    border-radius: 12px;
    margin-bottom: 20px;
    box-shadow: 0 6px 20px rgba(30, 60, 114, 0.4);
    border: 1px solid rgba(255,255,255,0.2);
}
#aegis_logo {
    text-align: center;
    margin-bottom: 20px;
    padding: 10px;
}
#aegis_logo img {
    max-width: 200px;
    height: auto;
    border-radius: 10px;
    box-shadow: 0 4px 15px rgba(30, 60, 114, 0.3);
    transition: all 0.3s ease;
}
#aegis_logo img:hover {
    transform: scale(1.05);
    box-shadow: 0 6px 20px rgba(30, 60, 114, 0.5);
}
.mission-section {
    background: linear-gradient(135deg, #2c1810 0%, #8b4513 50%, #d2691e 100%);
    padding: 25px;
    border-radius: 12px;
    margin: 20px 0;
    color: white;
    box-shadow: 0 8px 25px rgba(139, 69, 19, 0.4);
}
.investment-section {
    background: linear-gradient(135deg, #1a472a 0%, #2d5a3d 50%, #4a7c59 100%);
    padding: 25px;
    border-radius: 12px;
    margin: 20px 0;
    color: white;
    box-shadow: 0 8px 25px rgba(26, 71, 42, 0.4);
}
.quadrant-section {
    background: linear-gradient(135deg, #4a148c 0%, #7b1fa2 50%, #9c27b0 100%);
    padding: 25px;
    border-radius: 12px;
    margin: 20px 0;
    color: white;
    box-shadow: 0 8px 25px rgba(74, 20, 140, 0.4);
}
.contact-section {
    background: linear-gradient(135deg, #b71c1c 0%, #d32f2f 50%, #f44336 100%);
    padding: 20px;
    border-radius: 12px;
    margin: 20px 0;
    color: white;
    text-align: center;
    box-shadow: 0 8px 25px rgba(183, 28, 28, 0.4);
}
#generate_button {
    color: white;
    border-color: #1e3c72;
    background: linear-gradient(135deg, #1e3c72 0%, #2a5298 100%);
    width: 200px;
    height: 60px;
    border-radius: 10px;
    font-weight: bold;
    font-size: 16px;
    transition: all 0.3s ease;
    box-shadow: 0 4px 15px rgba(30, 60, 114, 0.3);
}
#generate_button:hover {
    transform: translateY(-3px);
    box-shadow: 0 8px 25px rgba(30, 60, 114, 0.5);
    background: linear-gradient(135deg, #2a5298 0%, #1e3c72 100%);
}
footer {
    visibility: hidden
}
body {
    background: linear-gradient(135deg, #0f2027 0%, #203a43 50%, #2c5364 100%);
    min-height: 100vh;
}
.gradio-container {
    background: linear-gradient(135deg, #0f2027 0%, #203a43 50%, #2c5364 100%);
    min-height: 100vh;
}
.gr-textbox {
    border: 2px solid #1e3c72;
    border-radius: 8px;
    background: rgba(255,255,255,0.95);
}
.gr-slider {
    accent-color: #1e3c72;
}
.gr-accordion {
    background: rgba(30, 60, 114, 0.1);
    border: 1px solid rgba(30, 60, 114, 0.3);
    border-radius: 8px;
}
"""


def _get_footer_message() -> str:
    version = f"<center><p> {APP_VERSION} "
    footer_msg = version + (
        '  © 2025 | Gaston Software Solutions Tec | '
        'Tel/WhatsApp: <a href="https://wa.me/256755274944" target="_blank" style="color: #90ee90;">+256755274944</a> | '
        'Email: <a href="mailto:gastonsoftwaresolution234@gmail.com" style="color: #90ee90;">gastonsoftwaresolution234@gmail.com</a> | '
        'AEGIS Bio-Digital Lab 10</p></center>'
    )
    warning_msg = """
    <div style="background: linear-gradient(135deg, #1e3c72 0%, #2a5298 100%); padding: 20px; border-radius: 10px; color: white; text-align: center; margin: 20px 0;">
        <p><b>🧬 AEGIS Bio-Digital Lab 10 - Advanced Bio-Digital Synthesis Platform</b></p>
        <p><strong>Powered by:</strong> Sovereign Quadrant AI — Synthesizing Reasoning, Genomics, Microscopy, and Volumetric Data</p>
        <p><strong>A Product of:</strong> Gaston Software Solutions Tec | <strong>Connect:</strong> Tel: +256755274944</p>
        <p><em>For research and simulation purposes.</em></p>
    </div><br>
    """
    
    # Green footer section
    green_footer = f"""
    <div style="background: linear-gradient(135deg, #1a472a 0%, #2d5a3d 50%, #4a7c59 100%); padding: 20px; border-radius: 12px; margin: 20px 0; color: white; text-align: center; box-shadow: 0 8px 25px rgba(26, 71, 42, 0.4);">
        <p style="font-size: 16px; font-weight: bold; margin-bottom: 10px; color: #90ee90;">{APP_VERSION}</p>
        <p style="font-size: 14px; margin-bottom: 8px;">© 2025 | <strong>Gaston Software Solutions Tec</strong></p>
        <p style="font-size: 14px; margin-bottom: 8px;">
            <strong>Tel/WhatsApp:</strong> <a href="https://wa.me/256755274944" target="_blank" style="color: #90ee90; text-decoration: none;">+256755274944</a> | 
            <strong>Email:</strong> <a href="mailto:gastonsoftwaresolution234@gmail.com" style="color: #90ee90; text-decoration: none;">gastonsoftwaresolution234@gmail.com</a>
        </p>
        <p style="font-size: 14px; color: #b8e6b8;"><strong>AEGIS Bio-Digital Lab 10</strong> - Building Humanity's Predictive Immune System</p>
    </div>
    """
    
    return warning_msg + green_footer


with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="container"):
        # AEGIS Logo at the top center
        gr.HTML(get_logo_html())
        
        use_openvino = "" if is_openvino_device() else ""
        gr.Markdown(
            f"""# 🧬 AEGIS Bio-Digital Lab 10 {use_openvino}
               **Artificially Expanded Genetic Information System**
               
               **Visualizing the Fabric of Life and Technology Through Time**
               
               *Advanced generative intelligence platform transforming complex scientific and technological concepts into high-fidelity visualizations*
               
               **Device:** {DEVICE.upper()} , {get_device_name()} | OpenVINO | **Powered by Sovereign Quadrant AI**
            """,
            elem_id="intro",
        )
        
        # Mission Statement: The Sovereign Timeline Project
        gr.HTML("""
        <div style="background: linear-gradient(135deg, #2c1810 0%, #8b4513 50%, #d2691e 100%); padding: 25px; border-radius: 12px; margin: 20px 0; color: white; box-shadow: 0 8px 25px rgba(139, 69, 19, 0.4);">
            <h2 style="text-align: center; margin-bottom: 20px; color: #ffd700;">🎯 Mission Statement: The Sovereign Timeline Project</h2>
            <p style="font-size: 18px; line-height: 1.6; margin-bottom: 15px; text-align: justify;">
                <strong>Our mission is to build a predictive immune system for humanity.</strong> The Artificially Expanded Genetic Information System (AEGIS) – Bio-Digital Lab 10 is not a tool for the present, but a shield for the future. We are developing the world's first temporal bio-digital modeling platform to visualize, understand, and preemptively counteract global biological threats before they emerge.
            </p>
            <p style="font-size: 16px; line-height: 1.6; text-align: justify;">
                By simulating pathogen evolution, mapping the hidden pathways of future pandemics, and designing proactive countermeasures across projected timelines, we aim to transform global health security from a reactive struggle into a predictive science.
            </p>
        </div>
        """)
        
        # Investment Call to Action
        gr.HTML("""
        <div style="background: linear-gradient(135deg, #1a472a 0%, #2d5a3d 50%, #4a7c59 100%); padding: 25px; border-radius: 12px; margin: 20px 0; color: white; box-shadow: 0 8px 25px rgba(26, 71, 42, 0.4);">
            <h2 style="text-align: center; margin-bottom: 20px; color: #90ee90;">📢 Investment Call: Partner in the Sovereign Timeline</h2>
            <p style="font-size: 18px; line-height: 1.6; margin-bottom: 15px; text-align: center;">
                <strong>The next great leap in human survival will not be fought in clinics alone, but in the sovereign digital space where time itself becomes a variable.</strong>
            </p>
            <div style="background: rgba(255,255,255,0.1); padding: 20px; border-radius: 8px; margin: 15px 0;">
                <h3 style="color: #90ee90; margin-bottom: 15px;">💰 Invest in the Predictive Frontier</h3>
                <p style="margin-bottom: 10px;">Your capital will directly fuel the core engines of the Sovereign Quadrant:</p>
                <ol style="margin: 10px 0; padding-left: 20px; line-height: 1.8;">
                    <li><strong>Expand the Timeline:</strong> Scale our high-fidelity models to simulate decades, not years, of viral and technological evolution.</li>
                    <li><strong>Fortify the Dataverse:</strong> Integrate next-generation genomic, climate, and epidemiological datasets to increase predictive accuracy.</li>
                    <li><strong>Democratize the Shield:</strong> Deploy public-facing simulation modules (like this Space) to accelerate global research and preparedness.</li>
                </ol>
            </div>
            <p style="font-size: 16px; line-height: 1.6; text-align: center; font-style: italic; margin: 15px 0;">
                "This is more than an investment in technology. It is an investment in time—buying humanity the critical hours, days, and years needed to prevent the next crisis."
            </p>
        </div>
        """)
        
        # Sovereign Quadrant Visual
        gr.HTML("""
        <div style="background: linear-gradient(135deg, #4a148c 0%, #7b1fa2 50%, #9c27b0 100%); padding: 25px; border-radius: 12px; margin: 20px 0; color: white; box-shadow: 0 8px 25px rgba(74, 20, 140, 0.4);">
            <h3 style="text-align: center; margin-bottom: 20px; color: #e1bee7;">🧭 The Sovereign Quadrant Engine</h3>
            <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 15px; margin: 20px 0;">
                <div style="background: rgba(255,255,255,0.15); padding: 20px; border-radius: 8px; text-align: center;">
                    <h4 style="color: #f3e5f5; margin-bottom: 10px;">🧠 REASONING</h4>
                    <p style="font-size: 14px;">Nemotron AI</p>
                </div>
                <div style="background: rgba(255,255,255,0.15); padding: 20px; border-radius: 8px; text-align: center;">
                    <h4 style="color: #f3e5f5; margin-bottom: 10px;">🧬 GENOMICS</h4>
                    <p style="font-size: 14px;">AEGIS Data</p>
                </div>
                <div style="background: rgba(255,255,255,0.15); padding: 20px; border-radius: 8px; text-align: center;">
                    <h4 style="color: #f3e5f5; margin-bottom: 10px;">🔬 MICROSCOPY</h4>
                    <p style="font-size: 14px;">2D/3D Imaging</p>
                </div>
                <div style="background: rgba(255,255,255,0.15); padding: 20px; border-radius: 8px; text-align: center;">
                    <h4 style="color: #f3e5f5; margin-bottom: 10px;">📊 VOLUME</h4>
                    <p style="font-size: 14px;">Temporal Models</p>
                </div>
            </div>
            <p style="text-align: center; font-style: italic; margin-top: 15px;">
                "Building the mirror that shows us our future, so we can change it."
            </p>
        </div>
        """)
        
        # Strategic Partnership Contact
        gr.HTML("""
        <div style="background: linear-gradient(135deg, #b71c1c 0%, #d32f2f 50%, #f44336 100%); padding: 20px; border-radius: 12px; margin: 20px 0; color: white; text-align: center; box-shadow: 0 8px 25px rgba(183, 28, 28, 0.4);">
            <h3 style="margin-bottom: 15px; color: #ffcdd2;">🤝 Strategic Investment Discussions</h3>
            <p style="font-size: 18px; font-weight: bold; margin-bottom: 10px;">Gaston Software Solutions Tec | Sovereign Timeline Division</p>
            <p style="font-size: 16px; margin-bottom: 5px;"><strong>Tel:</strong> +256755274944</p>
            <p style="font-size: 14px; font-style: italic;">Contact us to join the mission of building humanity's predictive immune system</p>
        </div>
        """)
        
        gr.HTML(
            f"""
            <div style="background: linear-gradient(135deg, #1e3c72 0%, #2a5298 100%); padding: 20px; border-radius: 10px; margin: 10px 0; color: white;">
                <h3>🔬 Bio-Digital Synthesis Modules</h3>
                <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(250px, 1fr)); gap: 15px; margin-top: 15px;">
                    <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 8px;">
                        <strong>🦠 Pathogen Timeline</strong><br>
                        <small>Models morphological evolution of viruses & bacteria across time</small>
                    </div>
                    <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 8px;">
                        <strong>🧬 Synthetic Genomics</strong><br>
                        <small>Renders novel genetic circuits and AEGIS base pair structures</small>
                    </div>
                    <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 8px;">
                        <strong>🔬 Multi-Scale Imaging</strong><br>
                        <small>From atomic protein folds to tissue-level pathology</small>
                    </div>
                    <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 8px;">
                        <strong>⚙️ Tech-Bio Fusion</strong><br>
                        <small>Bio-hybrid devices and computational biological systems</small>
                    </div>
                </div>
            </div>
            <p id="project-links" align="center">
                <a href='https://github.com/rupeshs/fastsdcpu'><img src='https://img.shields.io/badge/AEGIS-Bio Digital Lab-Purple'></a>
                <a href='#'><img src='https://img.shields.io/badge/Powered by-Sovereign Quadrant AI-Blue'></a>
                <a href='#'><img src='https://img.shields.io/badge/Gaston-Software Solutions-Green'></a>
            </p> 
                    """
        )

        with gr.Row():
            with gr.Column(scale=2):
                gr.HTML("""
                <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); padding: 15px; border-radius: 8px; color: white; margin-bottom: 15px;">
                    <h4>🧬 Bio-Digital Synthesis Instructions</h4>
                    <ol style="margin: 10px 0; padding-left: 20px;">
                        <li><strong>Select Temporal Domain:</strong> Past (historical pandemics), Present (current lab models), Future (speculative synthetic biology)</li>
                        <li><strong>Define Subject:</strong> Input target organism, mechanism, or bio-tech fusion concept</li>
                        <li><strong>Set Parameters:</strong> Adjust complexity, style (schematic vs. photorealistic), resolution</li>
                        <li><strong>Synthesize:</strong> Generate unique bio-digital visualization</li>
                    </ol>
                    <div style="background: rgba(255,255,255,0.1); padding: 10px; border-radius: 5px; margin-top: 10px;">
                        <strong>Example Prompts:</strong><br>
                        • "3D timeline showing influenza virus spike protein variation from 1918 to present"<br>
                        • "Glowing schematic view of expanded DNA helix with synthetic nucleotides"<br>
                        • "Cross-section from atomic detail of receptor to full cellular infection event"<br>
                        • "Neural interface device with organic circuitry integrated with semiconductor layers"
                    </div>
                </div>
                """)
            with gr.Row():
                prompt = gr.Textbox(
                    placeholder="Example: 'Generate a photorealistic cross-sectional view of a hypothesized Martian extremophile bacterium, based on terrestrial archaea genetics, with radiation-resistant protein capsid. Style: scientific illustration with labeled components.'",
                    scale=5,
                    container=False,
                    lines=3,
                    label="🧬 Bio-Digital Synthesis Prompt"
                )
                generate_btn = gr.Button(
                    "🔬 Synthesize",
                    scale=1,
                    elem_id="generate_button",
                )

        image = gr.Image(type="filepath", label="🧬 Bio-Digital Visualization Output")
        with gr.Accordion("🔬 Advanced Bio-Digital Parameters", open=False):
            gr.HTML("""
            <div style="background: #f0f8ff; padding: 10px; border-radius: 5px; margin-bottom: 10px;">
                <strong>Temporal Domains:</strong> Past (1800-1999) | Present (2000-2024) | Future (2025-2100)<br>
                <strong>Visualization Styles:</strong> Schematic, Photorealistic, Scientific Illustration, Cross-sectional
            </div>
            """)
            steps = gr.Slider(
                label="🔬 Synthesis Complexity",
                value=1,
                minimum=1,
                maximum=3,
                step=1,
                info="Higher values = more detailed bio-digital rendering"
            )
            seed = gr.Slider(
                randomize=True,
                minimum=0,
                maximum=999999999,
                label="🧬 Genetic Seed",
                step=1,
                info="Reproducible bio-digital synthesis"
            )
            seed_checkbox = gr.Checkbox(
                label="🔒 Lock Genetic Seed",
                value=False,
                interactive=True,
            )
        gr.HTML(_get_footer_message())

        inputs = [prompt, steps, seed, seed_checkbox]
        generate_btn.click(fn=predict, inputs=inputs, outputs=image)


def start_demo_text_to_image(share=False):
    demo.queue()
    demo.launch(share=share)