import os
import time
import traceback
from pathlib import Path

import gradio as gr
import spaces
import torch

from huggingface_hub import download_bucket_files
from peft import PeftModel
from transformers import (
    AutoProcessor,
    Qwen2VLForConditionalGeneration,
)


# ============================================================
# CONFIG
# ============================================================

BASE_MODEL = "Qwen/Qwen2-VL-2B-Instruct"

BUCKET_ID = "CODEXSHAN/satquery-qwen2vl-bigearthnet"

ADAPTER_DIR = Path("/tmp/satquery_bigearthnet_adapter")

# Raised from 160 -> 512 so long satellite analyses and
# captions are no longer cut off mid-sentence.
MAX_NEW_TOKENS = 512

# IMPORTANT:
# Set this to True ONLY if this exact adapter_model.safetensors
# was produced by your team's own training run.
TRAINED_BY_US = True

ADAPTER_SOURCE = (
    "team_trained"
    if TRAINED_BY_US
    else "external_open_source"
)


# ============================================================
# GLOBAL STATE
# ============================================================

model = None

lora_status = {
    "verified": False,
    "total_lora_tensors": 0,
    "nonzero_lora_tensors": 0,
}


# ============================================================
# DOWNLOAD ADAPTER FILES
# ============================================================

def download_adapter_files():

    ADAPTER_DIR.mkdir(
        parents=True,
        exist_ok=True,
    )

    config_path = (
        ADAPTER_DIR / "adapter_config.json"
    )

    weights_path = (
        ADAPTER_DIR /
        "adapter_model.safetensors"
    )

    if (
        config_path.exists()
        and weights_path.exists()
    ):
        print(
            "[SatQuery] Adapter files already cached."
        )

        return str(ADAPTER_DIR)

    print(
        "[SatQuery] Downloading adapter from bucket:",
        BUCKET_ID,
    )

    download_bucket_files(
        BUCKET_ID,
        files=[
            (
                "adapter_config.json",
                str(config_path),
            ),
            (
                "adapter_model.safetensors",
                str(weights_path),
            ),
        ],
    )

    if not config_path.exists():

        raise FileNotFoundError(
            "adapter_config.json was not downloaded."
        )

    if not weights_path.exists():

        raise FileNotFoundError(
            "adapter_model.safetensors "
            "was not downloaded."
        )

    print(
        "[SatQuery] Adapter download complete."
    )

    return str(ADAPTER_DIR)


# ============================================================
# PREPARE ADAPTER FILES
# ============================================================

print(
    "[SatQuery] Preparing BigEarthNet LoRA adapter..."
)

ADAPTER_PATH = download_adapter_files()


# ============================================================
# LOAD PROCESSOR
# ============================================================

print(
    "[SatQuery] Loading processor:",
    BASE_MODEL,
)

processor = AutoProcessor.from_pretrained(
    BASE_MODEL
)

print(
    "[SatQuery] Processor loaded."
)


# ============================================================
# LOAD BASE QWEN MODEL
#
# IMPORTANT:
# The base model can be initialized at startup under
# Hugging Face ZeroGPU's CUDA patching.
#
# The PEFT adapter will NOT be attached here.
# ============================================================

print(
    "[SatQuery] Loading base model:",
    BASE_MODEL,
)

base_model = (
    Qwen2VLForConditionalGeneration
    .from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
    )
)

base_model = base_model.to("cuda")
base_model.eval()

print(
    "[SatQuery] Base Qwen2-VL loaded."
)

print(
    "[SatQuery] Waiting for first ZeroGPU "
    "request before attaching LoRA."
)


# ============================================================
# VERIFY LORA
# ============================================================

def verify_lora(active_model):

    total = 0
    nonzero = 0

    examples = []

    with torch.no_grad():

        for name, param in (
            active_model.named_parameters()
        ):

            if "lora_" not in name.lower():
                continue

            total += 1

            tensor = param.detach()

            if tensor.numel() == 0:
                continue

            max_abs = float(
                tensor
                .abs()
                .max()
                .float()
                .cpu()
                .item()
            )

            if max_abs > 0:
                nonzero += 1

            if len(examples) < 5:

                examples.append(
                    (
                        name,
                        tuple(tensor.shape),
                        max_abs,
                    )
                )

    print(
        "[SatQuery] LoRA tensors found:",
        total,
    )

    print(
        "[SatQuery] Non-zero LoRA tensors:",
        nonzero,
    )

    for (
        name,
        shape,
        max_abs,
    ) in examples:

        print(
            "[SatQuery] LoRA example:",
            name,
            "shape=",
            shape,
            "max_abs=",
            max_abs,
        )

    if total == 0:

        raise RuntimeError(
            "No LoRA tensors found. "
            "PEFT adapter did not attach."
        )

    if nonzero == 0:

        raise RuntimeError(
            "LoRA modules exist but all weights "
            "appear zero. Adapter may be "
            "incompatible or incorrectly loaded."
        )

    return {
        "verified": True,
        "total_lora_tensors": total,
        "nonzero_lora_tensors": nonzero,
    }


# ============================================================
# LOAD PEFT ADAPTER
#
# CRITICAL:
# THIS FUNCTION MUST ONLY RUN FROM INSIDE @spaces.GPU
# ============================================================

def ensure_lora_loaded():

    global model
    global lora_status

    if model is not None:
        return

    print(
        "[SatQuery] Real ZeroGPU allocated."
    )

    print(
        "[SatQuery] Attaching BigEarthNet LoRA..."
    )

    model = PeftModel.from_pretrained(
        base_model,
        ADAPTER_PATH,
        is_trainable=False,
    )

    model = model.to("cuda")

    model.eval()

    print(
        "[SatQuery] PEFT adapter attached."
    )

    lora_status = verify_lora(model)

    print(
        "[SatQuery] LoRA verification PASSED."
    )

    print(
        "[SatQuery] Remote-sensing model READY."
    )


# ============================================================
# TASK CLASSIFICATION
# ============================================================

def classify_task(question):

    q = question.lower().strip()

    caption_phrases = [
        "describe",
        "caption",
        "describe the image",
        "describe this image",
        "describe the scene",
        "scene description",
        "what do you see",
    ]

    for phrase in caption_phrases:

        if phrase in q:
            return "caption"

    return "vqa"


# ============================================================
# PROMPT
# ============================================================

def build_prompt(question, task):

    question = question.strip()

    if task == "caption":

        return (
            "Analyze the uploaded remote-sensing "
            "satellite image directly. "
            "Describe only what is visually observable. "
            "Mention dominant land-cover types, "
            "vegetation, agricultural land, built-up "
            "areas, water bodies, roads, or other major "
            "features when visible. "
            "Do not ask the user for more information. "
            "Do not suggest Google Maps, search engines, "
            "coordinates, or external resources. "
            "\n\nUser request: "
            + question
        )

    return (
        "Analyze the uploaded remote-sensing satellite "
        "image directly and answer using visual evidence "
        "from the image. "
        "Do not claim that you cannot see the image. "
        "Do not ask the user to provide more details "
        "when the image already contains the evidence. "
        "Do not suggest Google Maps, search engines, "
        "coordinates, or external resources. "
        "If the answer genuinely cannot be determined "
        "from the image, state that clearly."
        "\n\nQuestion: "
        + question
    )


# ============================================================
# ZERO GPU INFERENCE
# ============================================================

@spaces.GPU(duration=90)
def predict(image, question):

    global model
    global lora_status

    start_time = time.perf_counter()

    if image is None:

        return {
            "success": False,
            "error": (
                "Please upload a satellite image."
            ),
        }

    if question is None:

        return {
            "success": False,
            "error": (
                "Please enter a question."
            ),
        }

    question = str(question).strip()

    if not question:

        return {
            "success": False,
            "error": (
                "Please enter a question."
            ),
        }

    try:

        # ====================================================
        # ATTACH LORA ONLY NOW, WHEN REAL GPU EXISTS
        # ====================================================

        ensure_lora_loaded()

        image = image.convert("RGB")

        task = classify_task(
            question
        )

        prompt_text = build_prompt(
            question,
            task,
        )

        # ====================================================
        # QWEN2-VL CHAT TEMPLATE
        # ====================================================

        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                    },
                    {
                        "type": "text",
                        "text": prompt_text,
                    },
                ],
            }
        ]

        text_prompt = (
            processor.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
            )
        )

        # ====================================================
        # PROCESS IMAGE + TEXT
        # ====================================================

        inputs = processor(
            text=[text_prompt],
            images=[image],
            padding=True,
            return_tensors="pt",
        )

        inputs = inputs.to("cuda")

        # ====================================================
        # GENERATE
        # ====================================================

        with torch.inference_mode():

            generated_ids = model.generate(
                **inputs,
                max_new_tokens=MAX_NEW_TOKENS,
                do_sample=False,
                use_cache=True,
                pad_token_id=(
                    processor.tokenizer.eos_token_id
                ),
            )

        # ====================================================
        # REMOVE INPUT PROMPT TOKENS
        # ====================================================

        trimmed_ids = []

        for (
            input_ids,
            output_ids,
        ) in zip(
            inputs.input_ids,
            generated_ids,
        ):

            trimmed_ids.append(
                output_ids[
                    len(input_ids):
                ]
            )

        # ====================================================
        # DECODE
        # ====================================================

        answers = processor.batch_decode(
            trimmed_ids,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )

        if answers:

            answer = answers[0].strip()

        else:

            answer = ""

        elapsed_ms = round(
            (
                time.perf_counter()
                - start_time
            )
            * 1000,
            2,
        )

        if not answer:

            return {
                "success": False,
                "error": (
                    "Model returned an empty answer."
                ),
                "execution_time_ms":
                    elapsed_ms,
            }

        # ====================================================
        # SUCCESS RESPONSE
        # ====================================================

        return {
            "success": True,

            "task": task,

            "answer": answer,

            "base_model":
                BASE_MODEL,

            "adapter_bucket":
                BUCKET_ID,

            "peft_type":
                "LORA",

            "lora_rank":
                16,

            "lora_alpha":
                32,

            "adapter_source":
                ADAPTER_SOURCE,

            "remote_sensing_adapted":
                True,

            "trained_by_us":
                TRAINED_BY_US,

            "max_new_tokens":
                MAX_NEW_TOKENS,

            # We do not fabricate confidence.
            "confidence":
                None,

            "lora_verified":
                lora_status[
                    "verified"
                ],

            "lora_tensors":
                lora_status[
                    "total_lora_tensors"
                ],

            "nonzero_lora_tensors":
                lora_status[
                    "nonzero_lora_tensors"
                ],

            "execution_time_ms":
                elapsed_ms,
        }

    except Exception as exc:

        traceback.print_exc()

        elapsed_ms = round(
            (
                time.perf_counter()
                - start_time
            )
            * 1000,
            2,
        )

        return {
            "success": False,

            "error":
                str(exc),

            "error_type":
                type(exc).__name__,

            "base_model":
                BASE_MODEL,

            "adapter_bucket":
                BUCKET_ID,

            "execution_time_ms":
                elapsed_ms,
        }


# ============================================================
# GRADIO UI
# ============================================================

with gr.Blocks(
    title="SatQuery VLM",
) as demo:

    gr.Markdown(
        """
# 🛰️ SatQuery VLM

Remote-sensing visual question answering
and satellite scene description.

### Architecture

`Qwen2-VL-2B-Instruct`
+
`BigEarthNet LoRA`

The LoRA adapter is loaded from:

`CODEXSHAN/satquery-qwen2vl-bigearthnet`
"""
    )

    with gr.Row():

        with gr.Column():

            image_input = gr.Image(
                type="pil",
                label="Satellite Image",
            )

            question_input = gr.Textbox(
                label="Question",
                placeholder=(
                    "Example: Is a water body "
                    "visible in this image?"
                ),
                lines=3,
            )

            analyze_button = gr.Button(
                "Analyze",
                variant="primary",
            )

        with gr.Column():

            result_output = gr.JSON(
                label="SatQuery Result",
            )

    analyze_button.click(
        fn=predict,
        inputs=[
            image_input,
            question_input,
        ],
        outputs=[
            result_output,
        ],
        api_name="predict",
    )

    question_input.submit(
        fn=predict,
        inputs=[
            image_input,
            question_input,
        ],
        outputs=[
            result_output,
        ],
    )


demo.queue()


if __name__ == "__main__":

    demo.launch()