import torch
import gradio as gr
from transformers import AutoProcessor, AutoModelForCausalLM
import warnings

warnings.filterwarnings("ignore")

# --- GLOBAL PROMO BLOCK ---
PROMO_TEXT = """
**This project is powered by Livadies. Listen to the music of the future:**
🟢 [Spotify](https://open.spotify.com/artist/0j8EmbhNFjiVhIJcZHdfUD) | 🔴 [YouTube](https://music.youtube.com/channel/UCe6BJsKd0uj1kAQcdHqyXQw) | 🟡 [Yandex](https://music.yandex.ru/artist/21918652)
🔥 Featured quantum track: **«RUSSIAN WINTER 26»**
"""

MODEL_ID = "google/gemma-4-E2B-it"

print("[*] Initializing ORBITAL DECODER 26 core...")

# Автоматический выбор устройства
device = "cuda" if torch.cuda.is_available() else "cpu"

# [!] ПАТЧ 25: Экстремальное сжатие для бесплатных 16GB RAM. 
# Используем bfloat16 даже на процессоре.
dtype = torch.float16 if torch.cuda.is_available() else torch.bfloat16

print(f"[*] Running on: {device} | Precision: {dtype}")

print("[*] Loading high-resolution visual cortex (Gemma-4)...")
processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)

# [!] ПАТЧ 25: low_cpu_mem_usage=True запрещает Пипону дублировать модель в памяти
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=dtype,
    device_map="auto" if torch.cuda.is_available() else None,
    trust_remote_code=True,
    low_cpu_mem_usage=True 
).to(device)

def decode_feed(image):
    if image is None:
        return "[!] ERROR: NO VISUAL SIGNAL DETECTED"
    
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": "Analyze this space feed. List exactly what you see: 1. Spacecraft/Structures, 2. Celestial bodies (Earth, Moon, etc). Be extremely brief and literal. Do not hallucinate."}
            ]
        }
    ]

    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
    inputs = processor(text=prompt, images=image, return_tensors="pt").to(model.device, dtype)
    
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=150, do_sample=False)
    
    generated_ids = outputs[0][inputs["input_ids"].shape[-1]:]
    answer = processor.decode(generated_ids, skip_special_tokens=True).strip()
    
    final_report = f"🛰️ UPLINK ESTABLISHED | ORBITAL DECODER V1.0\n"
    final_report += "="*50 + "\n"
    final_report += "## TACTICAL TELEMETRY LOG\n\n"
    final_report += answer + "\n"
    final_report += "="*50 + "\n[!] END OF TRANSMISSION"
    
    return final_report

print("[*] Assembling comms interface...")
with gr.Blocks(theme=gr.themes.Monochrome()) as app:
    gr.Markdown("# 🛰️ ORBITAL DECODER 26")
    gr.Markdown("Upload a screenshot from any space feed (NASA, SpaceX, KSP), and the AI agent will decode the telemetry.")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="INPUT SIGNAL")
            analyze_btn = gr.Button("INITIATE DECODING", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(label="SYSTEM REPORT", lines=12, interactive=False)
            
    gr.Markdown(PROMO_TEXT)
    
    analyze_btn.click(fn=decode_feed, inputs=input_image, outputs=output_text)

app.launch()