"""Hugging Face ZeroGPU Space: SatQuery TEOChat Official Bi-Temporal Reasoning Service.

Official Model: jirvin16/TEOChat
Official Repository: https://github.com/ermongroup/TEOChat
Responsibilities:
1. Bi-temporal satellite image reasoning (T1 -> T2)
2. Temporal change description
3. Change-based Visual Question Answering (VQA)
4. Temporal Earth Observation question answering

NO mock fallbacks. NO NumPy image differencing. NO hardcoded percentages or categories.
Uses real Video-LLaVA / TEOChat official inference code only.
"""
try:
    import spaces
    gpu_decorator = spaces.GPU(duration=20)
    has_spaces = True
except Exception:
    def gpu_decorator(fn):
        return fn
    has_spaces = False

import os
import io
import time
import base64
import tempfile
import logging
from typing import Optional, Dict, Any, List
from PIL import Image
import torch
import gradio as gr
from fastapi import FastAPI, Request, Response, status
from pydantic import BaseModel, Field

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("satquery.teochat")

MODEL_ID = os.getenv("MODEL_ID", "jirvin16/TEOChat")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

_model = None
_tokenizer = None
_processor = None
_load_error = None


def ensure_videollava_installed():
    """Ensure videollava package from official TEOChat repo is importable."""
    try:
        import videollava
        return
    except ImportError:
        logger.info("[TEOChat] videollava not found. Attempting to install from official repo...")
        import subprocess
        import sys
        try:
            subprocess.run(
                [sys.executable, "-m", "pip", "install", "--no-deps", "git+https://github.com/ermongroup/TEOChat.git"],
                check=True,
                capture_output=True
            )
            logger.info("[TEOChat] Successfully installed TEOChat videollava package!")
        except Exception as err:
            logger.error(f"[TEOChat] Failed to pip install git+https://github.com/ermongroup/TEOChat.git: {err}")
            raise ImportError(
                "Official TEOChat 'videollava' library is required. "
                "Please run: pip install git+https://github.com/ermongroup/TEOChat.git"
            ) from err


def load_teochat_model(target_device: Optional[str] = None):
    """Load the official TEOChat model using official videollava.eval.eval.load_model."""
    global _model, _tokenizer, _processor, _load_error
    resolved_dev = target_device or DEVICE
    if _model is not None and _tokenizer is not None and _processor is not None:
        if resolved_dev == "cpu":
            try:
                _model = _model.to("cpu")
            except Exception:
                pass
        return _tokenizer, _model, _processor

    logger.info(f"[TEOChat] Loading official model from {MODEL_ID} on {resolved_dev}...")
    try:
        ensure_videollava_installed()
        from videollava.eval.eval import load_model
        from videollava.eval.inference import run_inference_single

        # On ZeroGPU A10G (24GB VRAM), load in FP16 natively for maximum speed and stability
        use_8bit = False
        tokenizer, model, processor = load_model(
            model_path=MODEL_ID,
            model_base=None,
            load_8bit=use_8bit,
            device=resolved_dev
        )

        _tokenizer = tokenizer
        _model = model
        _processor = processor
        _load_error = None
        logger.info("[TEOChat] Successfully loaded official checkpoint!")
        return _tokenizer, _model, _processor
    except Exception as exc:
        _load_error = str(exc)
        logger.error(f"[TEOChat] Failed to load official model: {exc}", exc_info=True)
        raise exc


def decode_and_save_temp(img_input: Any, suffix: str = ".png") -> str:
    """Decode image payload and persist to a temporary file path for the official processor."""
    if img_input is None:
        raise ValueError("Missing required temporal image input.")

    # If it's already a valid existing file path on disk, return it directly
    if isinstance(img_input, str) and os.path.exists(img_input):
        return os.path.abspath(img_input)

    pil_img = None
    if isinstance(img_input, Image.Image):
        pil_img = img_input.convert("RGB")
    elif isinstance(img_input, str):
        if img_input.startswith("data:image"):
            _, b64 = img_input.split(",", 1)
            raw = base64.b64decode(b64)
            pil_img = Image.open(io.BytesIO(raw)).convert("RGB")
        else:
            try:
                raw = base64.b64decode(img_input)
                pil_img = Image.open(io.BytesIO(raw)).convert("RGB")
            except Exception as e:
                raise ValueError(f"Failed to decode base64 image: {e}")

    if pil_img is None:
        raise ValueError(f"Unrecognized image input format: {type(img_input)}")

    tmp_file = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
    pil_img.save(tmp_file.name, format="PNG")
    tmp_file.close()
    return tmp_file.name


def _run_teochat_inference_core(
    image_t1_input: Any,
    image_t2_input: Any,
    query: str,
    target_device: Optional[str] = None
) -> Dict[str, Any]:
    """Execute official TEOChat forward pass on bi-temporal sequence [T1, T2]."""
    t0 = time.time()
    tokenizer, model, processor = load_teochat_model(target_device=target_device)

    if target_device and str(next(model.parameters()).device) != target_device:
        try:
            model = model.to(target_device)
        except Exception as dev_err:
            logger.warning(f"[TEOChat] Could not move model to {target_device}: {dev_err}")

    if not query or not query.strip():
        query = "What changed between T1 and T2?"

    # 1. Save or expose valid local file paths
    t1_path = decode_and_save_temp(image_t1_input, suffix="_t1.png")
    t2_path = decode_and_save_temp(image_t2_input, suffix="_t2.png")

    try:
        # 2. Preserve strict chronological order [T1, T2]
        image_paths = [t1_path, t2_path]

        # 3. Build official prompt containing <video>
        inp = f"These are two satellite images in chronological order: <video> {query.strip()}"

        logger.info(f"[TEOChat] Executing official run_inference_single with {len(image_paths)} temporal images...")
        from videollava.eval.inference import run_inference_single
        with torch.inference_mode():
            response = run_inference_single(
                model=model,
                processor=processor,
                tokenizer=tokenizer,
                inp=inp,
                image_paths=image_paths,
                conv_mode="v1",
                timestamps=[],
                prompt_strategy="interleave",
                chronological_prefix=True,
                temperature=0.2,
                max_new_tokens=256
            )
    finally:
        # Clean up temporary files
        for p in [t1_path, t2_path]:
            try:
                if os.path.exists(p):
                    os.remove(p)
            except Exception:
                pass

    elapsed_sec = round(time.time() - t0, 3)

    # Return standard normalized payload
    return {
        "status": "success",
        "model": "TEOChat",
        "task": "temporal_change_qa",
        "answer": response,
        "confidence": None,  # No fabricated confidence score
        "metadata": {
            "model_id": MODEL_ID,
            "checkpoint": "jirvin16/TEOChat",
            "num_temporal_images": 2,
            "inference_time_sec": elapsed_sec,
            "device": str(next(model.parameters()).device)
        }
    }


@gpu_decorator
def _run_teochat_inference_gpu(t1_img, t2_img, query):
    return _run_teochat_inference_core(t1_img, t2_img, query, target_device="cuda")


def _run_teochat_inference_cpu(t1_img, t2_img, query):
    return _run_teochat_inference_core(t1_img, t2_img, query, target_device="cpu")


# ==============================================================================
# Gradio Web Interface
# ==============================================================================
def gradio_fn(t1_img, t2_img, query):
    try:
        res = _run_teochat_inference_gpu(t1_img, t2_img, query)
        return res["answer"]
    except Exception as err:
        err_msg = str(err)
        logger.error(f"[TEOChat] GPU inference failure: {err_msg}", exc_info=True)
        logger.warning("[TEOChat] GPU execution encountered an issue. Seamlessly falling back to CPU with official model...")
        try:
            res = _run_teochat_inference_cpu(t1_img, t2_img, query)
            return res["answer"]
        except Exception as cpu_err:
            logger.error(f"[TEOChat] CPU fallback failure: {cpu_err}", exc_info=True)
            return f"Error: {str(cpu_err)}"


demo = gr.Interface(
    fn=gradio_fn,
    inputs=[
        gr.Image(type="pil", label="Baseline Observation (T1 / Pre-Event)"),
        gr.Image(type="pil", label="Target Observation (T2 / Post-Event)"),
        gr.Textbox(lines=2, placeholder="e.g. What changed between T1 and T2?", label="Temporal Reasoning Query", value="What changed between T1 and T2?")
    ],
    outputs=[
        gr.Textbox(label="TEOChat Official Delta Narrative")
    ],
    title="🛰️ SatQuery: Official TEOChat Bi-Temporal Specialist Service",
    description="Official jirvin16/TEOChat Bi-Temporal Remote Sensing Specialist. Temporal change description and change-based VQA.",
    api_name="predict"
)


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)