import asyncio
import base64
import io
import json
import logging
import os
import subprocess
import tempfile
import time
import uuid
from typing import Any

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image


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

HOST = os.getenv("HOST", "0.0.0.0")
PORT = int(os.getenv("PORT", "7860"))

MODEL_PATH = os.getenv(
    "LOCATEANYTHING_MODEL",
    "/models/locate-anything-q4_k.gguf",
)

CLI_PATH = os.getenv(
    "LOCATEANYTHING_CLI",
    "/usr/local/bin/locate-anything-cli",
)

CPU_THREADS = int(
    os.getenv("OMP_NUM_THREADS", "2")
)

MAX_IMAGE_SIZE = int(
    os.getenv("MAX_IMAGE_SIZE", "768")
)

INFERENCE_TIMEOUT = int(
    os.getenv("INFERENCE_TIMEOUT", "240")
)


# ============================================================
# LOGGING
# ============================================================

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
)

logger = logging.getLogger(
    "locateanything-live"
)


# ============================================================
# APP
# ============================================================

app = FastAPI(
    title="LocateAnything Live V2",
    version="2.1.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ============================================================
# STATIC
# ============================================================

if os.path.isdir("static"):
    app.mount(
        "/static",
        StaticFiles(directory="static"),
        name="static",
    )


# ============================================================
# STATE
# ============================================================

model_ready = False

# Only one CPU inference at once.
inference_lock = asyncio.Lock()

# Background jobs.
jobs = {}

# Most recent completed result.
last_result = {
    "status": "none",
    "job_id": None,
    "detections": [],
    "latency": None,
    "query": None,
    "width": None,
    "height": None,
}


# ============================================================
# HELPERS
# ============================================================

def executable_exists():
    return (
        os.path.isfile(CLI_PATH)
        and os.access(CLI_PATH, os.X_OK)
    )


def model_exists():
    return os.path.isfile(MODEL_PATH)


def model_size_gb():
    try:
        return round(
            os.path.getsize(MODEL_PATH)
            / (1024 ** 3),
            2,
        )
    except Exception:
        return None


def safe_number(
    value,
    default=None,
):
    try:
        if value is None:
            return default

        return float(value)

    except (
        TypeError,
        ValueError,
    ):
        return default


# ============================================================
# MODEL
# ============================================================

def initialize_model():

    global model_ready

    if model_ready:
        return

    logger.info(
        "Checking LocateAnything CPU runtime..."
    )

    logger.info(
        "CLI: %s",
        CLI_PATH,
    )

    logger.info(
        "Model: %s",
        MODEL_PATH,
    )

    logger.info(
        "Model size: %s GB",
        model_size_gb(),
    )

    if not executable_exists():
        raise RuntimeError(
            f"LocateAnything CLI not found: {CLI_PATH}"
        )

    if not model_exists():
        raise RuntimeError(
            f"LocateAnything model not found: {MODEL_PATH}"
        )

    process = subprocess.run(
        [
            CLI_PATH,
            "info",
            "--model",
            MODEL_PATH,
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        encoding="utf-8",
        errors="replace",
        timeout=60,
    )

    if process.returncode != 0:
        raise RuntimeError(
            process.stderr
            or process.stdout
            or "LocateAnything model check failed."
        )

    model_ready = True

    logger.info(
        "LocateAnything CPU runtime is ready."
    )


# ============================================================
# IMAGE
# ============================================================

def decode_data_url(
    data_url: str,
):

    if not data_url:
        raise ValueError(
            "Image data is empty."
        )

    if "," in data_url:
        data_url = data_url.split(
            ",",
            1,
        )[1]

    try:
        raw = base64.b64decode(
            data_url,
            validate=False,
        )
    except Exception as exc:
        raise ValueError(
            "Invalid base64 image."
        ) from exc

    try:
        image = Image.open(
            io.BytesIO(raw)
        )

        image.load()

    except Exception as exc:
        raise ValueError(
            "Unable to decode image."
        ) from exc

    return image.convert("RGB")


def prepare_image(
    image: Image.Image,
):

    width, height = image.size

    largest = max(
        width,
        height,
    )

    if largest <= MAX_IMAGE_SIZE:
        return image

    ratio = (
        MAX_IMAGE_SIZE
        / float(largest)
    )

    new_width = max(
        1,
        round(width * ratio),
    )

    new_height = max(
        1,
        round(height * ratio),
    )

    return image.resize(
        (
            new_width,
            new_height,
        ),
        Image.Resampling.LANCZOS,
    )


def save_image(
    image: Image.Image,
):

    file = tempfile.NamedTemporaryFile(
        suffix=".jpg",
        delete=False,
    )

    path = file.name

    file.close()

    image.save(
        path,
        format="JPEG",
        quality=82,
        optimize=True,
    )

    return path


# ============================================================
# CLI
# ============================================================

def build_command(
    image_path,
    query,
    mode,
    output_path,
):

    query = (
        query
        or "all objects"
    ).strip()[:500]

    mode = (
        mode
        or "hybrid"
    ).strip().lower()

    if mode not in {
        "hybrid",
        "slow",
        "fast",
    }:
        mode = "hybrid"

    return [
        CLI_PATH,

        "detect",

        "--model",
        MODEL_PATH,

        "--input",
        image_path,

        "--prompt",
        query,

        "--output",
        output_path,

        "--mode",
        mode,

        "--threads",
        str(CPU_THREADS),
    ]


# ============================================================
# JSON
# ============================================================

def load_json_file(path):

    if not os.path.isfile(path):
        return None

    try:

        with open(
            path,
            "r",
            encoding="utf-8",
        ) as file:

            text = file.read().strip()

        if not text:
            return None

        return json.loads(text)

    except Exception as exc:

        logger.warning(
            "Could not parse JSON: %s",
            exc,
        )

        return None


def parse_stdout_json(stdout):

    stdout = (
        stdout or ""
    ).strip()

    if not stdout:
        return None

    try:
        return json.loads(stdout)
    except json.JSONDecodeError:
        pass

    start = stdout.find("{")
    end = stdout.rfind("}")

    if (
        start >= 0
        and end > start
    ):

        try:
            return json.loads(
                stdout[start:end + 1]
            )
        except Exception:
            pass

    start = stdout.find("[")
    end = stdout.rfind("]")

    if (
        start >= 0
        and end > start
    ):

        try:
            return json.loads(
                stdout[start:end + 1]
            )
        except Exception:
            pass

    return None


# ============================================================
# DETECTION PARSER
# ============================================================

def extract_box(value):

    if isinstance(
        value,
        (list, tuple),
    ):

        if len(value) >= 4:

            values = [
                safe_number(value[0]),
                safe_number(value[1]),
                safe_number(value[2]),
                safe_number(value[3]),
            ]

            if all(
                x is not None
                for x in values
            ):
                return values

    if not isinstance(
        value,
        dict,
    ):
        return None

    for key in (
        "box",
        "bbox",
        "bounding_box",
        "boundingBox",
        "coordinates",
        "rect",
    ):

        if key in value:

            result = extract_box(
                value[key]
            )

            if result:
                return result

    x1 = value.get("x1")
    y1 = value.get("y1")
    x2 = value.get("x2")
    y2 = value.get("y2")

    if all(
        x is not None
        for x in (
            x1,
            y1,
            x2,
            y2,
        )
    ):

        return [
            safe_number(x1),
            safe_number(y1),
            safe_number(x2),
            safe_number(y2),
        ]

    xmin = value.get("xmin")
    ymin = value.get("ymin")
    xmax = value.get("xmax")
    ymax = value.get("ymax")

    if all(
        x is not None
        for x in (
            xmin,
            ymin,
            xmax,
            ymax,
        )
    ):

        return [
            safe_number(xmin),
            safe_number(ymin),
            safe_number(xmax),
            safe_number(ymax),
        ]

    return None


def extract_label(item):

    for key in (
        "label",
        "text",
        "name",
        "class",
        "object",
        "category",
        "description",
    ):

        value = item.get(key)

        if value is not None:
            return str(value)[:120]

    return "object"


def extract_score(item):

    for key in (
        "score",
        "confidence",
        "probability",
        "confidence_score",
    ):

        value = item.get(key)

        if value is not None:

            result = safe_number(
                value
            )

            if result is not None:
                return result

    return None


def normalize_box(
    box,
    width,
    height,
):

    if not box:
        return None

    x1, y1, x2, y2 = box

    if any(
        x is None
        for x in box
    ):
        return None

    maximum = max(
        abs(x1),
        abs(y1),
        abs(x2),
        abs(y2),
    )

    # 0..1
    if maximum <= 1.5:

        x1 *= width
        x2 *= width

        y1 *= height
        y2 *= height

    # 0..1000
    elif (
        maximum <= 1000
        and (
            x2 > width
            or y2 > height
        )
    ):

        x1 = x1 / 1000 * width
        x2 = x2 / 1000 * width

        y1 = y1 / 1000 * height
        y2 = y2 / 1000 * height

    x1 = max(
        0,
        min(width, x1),
    )

    x2 = max(
        0,
        min(width, x2),
    )

    y1 = max(
        0,
        min(height, y1),
    )

    y2 = max(
        0,
        min(height, y2),
    )

    if x2 < x1:
        x1, x2 = x2, x1

    if y2 < y1:
        y1, y2 = y2, y1

    if (
        x2 - x1 < 2
        or y2 - y1 < 2
    ):
        return None

    return {
        "x1": x1,
        "y1": y1,
        "x2": x2,
        "y2": y2,
    }


def find_detection_lists(data):

    if isinstance(
        data,
        list,
    ):
        return [data]

    if not isinstance(
        data,
        dict,
    ):
        return []

    results = []

    for key in (
        "detections",
        "boxes",
        "results",
        "objects",
        "predictions",
        "instances",
        "locations",
        "items",
    ):

        value = data.get(key)

        if isinstance(
            value,
            list,
        ):
            results.append(value)

    for value in data.values():

        if isinstance(
            value,
            dict,
        ):

            results.extend(
                find_detection_lists(
                    value
                )
            )

    return results


def normalize_detections(
    data,
    width,
    height,
):

    detections = []

    for items in find_detection_lists(
        data
    ):

        for item in items:

            if not isinstance(
                item,
                dict,
            ):
                continue

            box = extract_box(
                item
            )

            if not box:
                continue

            box = normalize_box(
                box,
                width,
                height,
            )

            if not box:
                continue

            detections.append(
                {
                    **box,

                    "label":
                        extract_label(
                            item
                        ),

                    "score":
                        extract_score(
                            item
                        ),
                }
            )

    # Remove duplicates.

    unique = []
    seen = set()

    for item in detections:

        key = (
            item["label"],
            round(item["x1"], 1),
            round(item["y1"], 1),
            round(item["x2"], 1),
            round(item["y2"], 1),
        )

        if key in seen:
            continue

        seen.add(key)
        unique.append(item)

    return unique


# ============================================================
# SYNCHRONOUS CLI WORKER
# ============================================================

def run_inference_sync(
    image_path,
    query,
    mode,
    width,
    height,
):

    output_file = tempfile.NamedTemporaryFile(
        suffix=".json",
        delete=False,
    )

    output_path = output_file.name

    output_file.close()

    command = build_command(
        image_path,
        query,
        mode,
        output_path,
    )

    logger.info(
        "Running LocateAnything:"
    )

    logger.info(
        "%s",
        " ".join(command),
    )

    environment = os.environ.copy()

    environment[
        "OMP_NUM_THREADS"
    ] = str(
        CPU_THREADS
    )

    environment[
        "GGML_NUM_THREADS"
    ] = str(
        CPU_THREADS
    )

    started = time.perf_counter()

    try:

        process = subprocess.run(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            encoding="utf-8",
            errors="replace",
            env=environment,
            timeout=INFERENCE_TIMEOUT,
        )

    except subprocess.TimeoutExpired:

        raise RuntimeError(
            "LocateAnything inference timed out."
        )

    finally:

        pass

    latency = round(
        (
            time.perf_counter()
            - started
        ) * 1000
    )

    stdout = (
        process.stdout or ""
    ).strip()

    stderr = (
        process.stderr or ""
    ).strip()

    logger.info(
        "Inference finished in %.2f seconds",
        latency / 1000,
    )

    if stderr:

        logger.info(
            "LocateAnything stderr:\n%s",
            stderr[-5000:],
        )

    if process.returncode != 0:

        try:
            os.remove(output_path)
        except OSError:
            pass

        raise RuntimeError(
            "LocateAnything exited with "
            f"code {process.returncode}.\n"
            + (
                stderr
                or stdout
                or "Unknown error."
            )
        )

    data = load_json_file(
        output_path
    )

    if data is None:

        data = parse_stdout_json(
            stdout
        )

    if data is None:

        data = {
            "detections": []
        }

    logger.info(
        "========== LOCATEANYTHING JSON RESULT =========="
    )

    logger.info(
        "%s",
        json.dumps(
            data,
            indent=2,
            ensure_ascii=False,
        )[:15000],
    )

    logger.info(
        "================================================="
    )

    detections = normalize_detections(
        data,
        width,
        height,
    )

    try:
        os.remove(output_path)
    except OSError:
        pass

    return {
        "status": "success",

        "detections":
            detections,

        "raw":
            data,

        "stdout":
            stdout,

        "stderr":
            stderr,

        "latency":
            latency,

        "width":
            width,

        "height":
            height,

        "finished_at":
            time.time(),
    }


# ============================================================
# BACKGROUND JOB
# ============================================================

async def inference_job(
    job_id,
    image_path,
    query,
    mode,
    width,
    height,
):

    job = jobs.get(job_id)

    if not job:
        return

    job["status"] = "running"

    job["started_at"] = time.time()

    try:

        initialize_model()

        async with inference_lock:

            result = await asyncio.to_thread(
                run_inference_sync,

                image_path,

                query,

                mode,

                width,

                height,
            )

        job.update(
            result
        )

        job["status"] = "completed"

        global last_result

        last_result = {
            **result,

            "job_id":
                job_id,

            "query":
                query,

            "status":
                "completed",
        }

        logger.info(
            "Job %s completed: %d detections",
            job_id,
            len(
                result["detections"]
            ),
        )

    except Exception as exc:

        logger.exception(
            "Job %s failed",
            job_id,
        )

        job["status"] = "error"

        job["error"] = (
            type(exc).__name__
        )

        job["message"] = str(exc)

    finally:

        try:
            os.remove(image_path)
        except OSError:
            pass


# ============================================================
# CREATE JOB
# ============================================================

async def create_job(
    image,
    query,
    mode,
):

    image = prepare_image(
        image
    )

    image_path = save_image(
        image
    )

    job_id = uuid.uuid4().hex

    jobs[job_id] = {
        "job_id":
            job_id,

        "status":
            "queued",

        "query":
            query,

        "mode":
            mode,

        "width":
            image.width,

        "height":
            image.height,

        "created_at":
            time.time(),
    }

    asyncio.create_task(
        inference_job(
            job_id,
            image_path,
            query,
            mode,
            image.width,
            image.height,
        )
    )

    return jobs[job_id]


# ============================================================
# FRONTEND
# ============================================================

@app.get("/")
async def index():

    path = "static/index.html"

    if os.path.isfile(path):

        return FileResponse(
            path
        )

    return {
        "name":
            "LocateAnything Live V2",

        "status":
            "online",
    }


# ============================================================
# API
# ============================================================

@app.get("/api")
async def api():

    running = sum(
        1
        for job in jobs.values()
        if job.get("status")
        in {
            "queued",
            "running",
        }
    )

    return {
        "name":
            "LocateAnything Live V2",

        "status":
            "online",

        "runtime":
            "locate-anything.cpp",

        "model":
            MODEL_PATH,

        "device":
            "cpu",

        "cuda":
            False,

        "model_loaded":
            model_ready,

        "model_exists":
            model_exists(),

        "cli_exists":
            executable_exists(),

        "model_size_gb":
            model_size_gb(),

        "cpu_threads":
            CPU_THREADS,

        "running_jobs":
            running,
    }


# ============================================================
# HEALTH
# ============================================================

@app.get("/health")
async def health():

    return {
        "status":
            "healthy",

        "runtime":
            "locate-anything.cpp",

        "device":
            "cpu",

        "cuda":
            False,

        "model_loaded":
            model_ready,

        "model_exists":
            model_exists(),

        "cli_exists":
            executable_exists(),
    }


# ============================================================
# WARMUP
# ============================================================

@app.post("/warmup")
async def warmup():

    try:

        started = time.perf_counter()

        initialize_model()

        return {
            "status":
                "ready",

            "model_loaded":
                True,

            "device":
                "cpu",

            "cuda":
                False,

            "time_ms":
                round(
                    (
                        time.perf_counter()
                        - started
                    ) * 1000
                ),
        }

    except Exception as exc:

        return {
            "status":
                "error",

            "model_loaded":
                False,

            "device":
                "cpu",

            "cuda":
                False,

            "error":
                type(exc).__name__,

            "message":
                str(exc),
        }


# ============================================================
# JOB API
# ============================================================

@app.get("/job/{job_id}")
async def get_job(
    job_id: str,
):

    job = jobs.get(
        job_id
    )

    if not job:

        return {
            "status":
                "not_found",

            "job_id":
                job_id,
        }

    return job


# ============================================================
# LAST RESULT
# ============================================================

@app.get("/last-result")
async def get_last_result():

    return last_result


# ============================================================
# RUNTIME
# ============================================================

@app.get("/runtime")
async def runtime():

    if not executable_exists():

        return {
            "status":
                "cli_missing",

            "cli":
                CLI_PATH,
        }

    try:

        process = subprocess.run(
            [
                CLI_PATH,
                "--help",
            ],

            stdout=subprocess.PIPE,

            stderr=subprocess.PIPE,

            text=True,

            encoding="utf-8",

            errors="replace",

            timeout=15,
        )

        return {
            "status":
                "ok",

            "cli":
                CLI_PATH,

            "model":
                MODEL_PATH,

            "cli_exists":
                True,

            "model_exists":
                model_exists(),

            "model_size_gb":
                model_size_gb(),

            "stdout":
                process.stdout,

            "stderr":
                process.stderr,
        }

    except Exception as exc:

        return {
            "status":
                "error",

            "error":
                type(exc).__name__,

            "message":
                str(exc),
        }


# ============================================================
# WEBSOCKET
# ============================================================

@app.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
):

    await websocket.accept()

    logger.info(
        "WebSocket client connected."
    )

    try:

        await websocket.send_json(
            {
                "type":
                    "connected",

                "runtime":
                    "locate-anything.cpp",

                "device":
                    "cpu",

                "cuda":
                    False,

                "model":
                    MODEL_PATH,

                "model_loaded":
                    model_ready,
            }
        )

        while True:

            raw = (
                await websocket.receive_text()
            )

            try:

                payload = json.loads(
                    raw
                )

            except json.JSONDecodeError:

                await websocket.send_json(
                    {
                        "type":
                            "error",

                        "message":
                            "Invalid JSON.",
                    }
                )

                continue

            message_type = payload.get(
                "type"
            )


            # ------------------------------------------------
            # PING
            # ------------------------------------------------

            if message_type == "ping":

                await websocket.send_json(
                    {
                        "type":
                            "pong",

                        "time":
                            time.time(),
                    }
                )

                continue


            # ------------------------------------------------
            # STOP
            # ------------------------------------------------

            if message_type == "stop":

                await websocket.send_json(
                    {
                        "type":
                            "stopped",
                    }
                )

                continue


            # ------------------------------------------------
            # ASYNC FRAME
            # ------------------------------------------------

            if message_type == "frame":

                image_data = payload.get(
                    "image"
                )

                query = (
                    payload.get(
                        "query",
                        "all objects",
                    )
                    or "all objects"
                )

                mode = (
                    payload.get(
                        "mode",
                        "hybrid",
                    )
                    or "hybrid"
                )

                frame_id = payload.get(
                    "frame_id"
                )


                if not image_data:

                    await websocket.send_json(
                        {
                            "type":
                                "error",

                            "message":
                                "Missing image.",

                            "frame_id":
                                frame_id,
                        }
                    )

                    continue


                try:

                    image = decode_data_url(
                        image_data
                    )

                except Exception as exc:

                    await websocket.send_json(
                        {
                            "type":
                                "error",

                            "message":
                                "Invalid image.",

                            "detail":
                                str(exc),

                            "frame_id":
                                frame_id,
                        }
                    )

                    continue


                # --------------------------------------------
                # Don't queue unlimited CPU jobs.
                # --------------------------------------------

                busy = any(
                    job.get("status")
                    in {
                        "queued",
                        "running",
                    }
                    for job in jobs.values()
                )


                if busy:

                    await websocket.send_json(
                        {
                            "type":
                                "busy",

                            "frame_id":
                                frame_id,

                            "message":
                                "Previous CPU inference is still running.",
                        }
                    )

                    continue


                job = await create_job(
                    image,
                    query,
                    mode,
                )


                await websocket.send_json(
                    {
                        "type":
                            "job",

                        "job_id":
                            job["job_id"],

                        "frame_id":
                            frame_id,

                        "status":
                            "queued",

                        "query":
                            query,

                        "message":
                            "Frame queued for CPU inference.",
                    }
                )

                continue


            # ------------------------------------------------
            # UNKNOWN
            # ------------------------------------------------

            await websocket.send_json(
                {
                    "type":
                        "error",

                    "message":
                        "Unknown message type.",
                }
            )


    except WebSocketDisconnect:

        logger.info(
            "WebSocket client disconnected."
        )

    except Exception as exc:

        logger.exception(
            "WebSocket error: %s",
            exc,
        )


# ============================================================
# START
# ============================================================

if __name__ == "__main__":

    import uvicorn

    logger.info(
        "=========================================="
    )

    logger.info(
        "LocateAnything Live V2"
    )

    logger.info(
        "Runtime: locate-anything.cpp"
    )

    logger.info(
        "Device: CPU"
    )

    logger.info(
        "Model: %s",
        MODEL_PATH,
    )

    logger.info(
        "CLI: %s",
        CLI_PATH,
    )

    logger.info(
        "CPU threads: %s",
        CPU_THREADS,
    )

    logger.info(
        "=========================================="
    )

    uvicorn.run(
        app,
        host=HOST,
        port=PORT,
        log_level="info",
    )