from __future__ import annotations

import argparse
import gc
import io
import json
import os
import sys
import tempfile
import threading
import time
import uuid
import base64
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path


MODEL = None
TOKENIZER = None
PROCESSOR = None
EDIT_MODEL = None
SERVER_CONFIG: dict = {}
LOAD_ERROR: str | None = None
IMAGE_JOBS: dict[str, dict] = {}
IMAGE_JOB_LOCK = threading.Lock()
IMAGE_JOB_CANCEL_EVENTS: dict[str, threading.Event] = {}
IMAGE_MODEL_LOCK = threading.Lock()
IMAGE_EDIT_MODEL_LOCK = threading.Lock()
MODEL_LOAD_STATE: dict[str, dict] = {
    "image": {"status": "not_loaded"},
    "image_edit": {"status": "not_loaded"},
}


def set_load_state(key: str, **fields: object) -> None:
    state = MODEL_LOAD_STATE.setdefault(key, {})
    state.update(fields)
    state["updated_at"] = time.time()


def release_cuda_memory() -> None:
    gc.collect()
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
    except Exception:
        return


def unload_image_model() -> None:
    global MODEL
    if MODEL is None:
        return
    MODEL = None
    release_cuda_memory()
    set_load_state("image", status="not_loaded", unloaded_at=time.time())


def unload_image_edit_model() -> None:
    global EDIT_MODEL
    if EDIT_MODEL is None:
        return
    EDIT_MODEL = None
    release_cuda_memory()
    set_load_state("image_edit", status="not_loaded", unloaded_at=time.time())


def has_active_image_job(operation: str) -> bool:
    active_statuses = {"queued", "loading_model", "running", "cancel_requested"}
    with IMAGE_JOB_LOCK:
        return any(
            job.get("operation") == operation and job.get("status") in active_statuses
            for job in IMAGE_JOBS.values()
        )


def diffusers_class_name(model_path: str | None) -> str:
    if not model_path:
        return ""
    index_path = Path(model_path) / "model_index.json"
    if not index_path.exists():
        return ""
    try:
        data = json.loads(index_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return ""
    return str(data.get("_class_name") or "")


def is_qwen_image_class(class_name: str) -> bool:
    return class_name in {"QwenImagePipeline", "QwenImageEditPipeline"}


def has_fp16_variant(model_path: str) -> bool:
    path = Path(model_path)
    if not path.exists():
        return False
    return any(path.rglob("*fp16.safetensors"))


def prepare_diffusers_pipeline(pipe, *, torch_module, class_name: str) -> None:
    if hasattr(pipe, "enable_attention_slicing"):
        pipe.enable_attention_slicing()
    if not torch_module.cuda.is_available():
        return
    if is_qwen_image_class(class_name) and hasattr(pipe, "enable_model_cpu_offload"):
        pipe.enable_model_cpu_offload()
    else:
        pipe.to("cuda")


def load_diffusers_pipeline(model_path: str, *, edit: bool):
    import torch

    class_name = diffusers_class_name(model_path)
    dtype = torch.bfloat16 if is_qwen_image_class(class_name) and torch.cuda.is_available() else torch.float16 if torch.cuda.is_available() else torch.float32
    variant = "fp16" if torch.cuda.is_available() and has_fp16_variant(model_path) else None
    if class_name == "QwenImageEditPipeline":
        from diffusers import QwenImageEditPipeline

        pipe = QwenImageEditPipeline.from_pretrained(model_path, torch_dtype=dtype)
    elif class_name == "QwenImagePipeline":
        from diffusers import QwenImagePipeline

        pipe = QwenImagePipeline.from_pretrained(model_path, torch_dtype=dtype)
    elif edit:
        from diffusers import StableDiffusionXLInpaintPipeline

        pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
            model_path,
            torch_dtype=dtype,
            use_safetensors=True,
            variant=variant,
        )
    else:
        from diffusers import StableDiffusionXLPipeline

        pipe = StableDiffusionXLPipeline.from_pretrained(
            model_path,
            torch_dtype=dtype,
            use_safetensors=True,
            variant=variant,
        )
    return pipe, torch, class_name


def json_response(handler: BaseHTTPRequestHandler, status: int, payload: dict) -> None:
    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    handler.send_response(status)
    handler.send_header("Content-Type", "application/json; charset=utf-8")
    handler.send_header("Content-Length", str(len(body)))
    handler.end_headers()
    handler.wfile.write(body)


def read_body(handler: BaseHTTPRequestHandler) -> dict:
    length = int(handler.headers.get("Content-Length") or "0")
    if length <= 0:
        return {}
    return json.loads(handler.rfile.read(length).decode("utf-8"))


def last_user_message(messages: list[dict]) -> str:
    for message in reversed(messages):
        if message.get("role") == "user":
            content = message.get("content", "")
            if isinstance(content, str):
                return content
            if isinstance(content, list):
                parts = [part.get("text", "") for part in content if isinstance(part, dict)]
                return "\n".join(part for part in parts if part)
    return ""


def load_translation_model() -> None:
    global LOAD_ERROR, MODEL, TOKENIZER
    if MODEL is not None and TOKENIZER is not None:
        return
    try:
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        model_path = SERVER_CONFIG["model_path"]
        TOKENIZER = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        MODEL = AutoModelForCausalLM.from_pretrained(
            model_path,
            device_map="auto",
            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else "auto",
            trust_remote_code=True,
        )
        LOAD_ERROR = None
    except Exception as exc:  # pragma: no cover - depends on host GPU/runtime
        LOAD_ERROR = f"{type(exc).__name__}: {exc}"
        raise


def load_audio_model() -> None:
    global LOAD_ERROR, MODEL, PROCESSOR
    if MODEL is not None and PROCESSOR is not None:
        return
    try:
        import torch
        from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline

        model_path = SERVER_CONFIG["model_path"]
        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
        MODEL = AutoModelForSpeechSeq2Seq.from_pretrained(
            model_path,
            torch_dtype=dtype,
            low_cpu_mem_usage=True,
            use_safetensors=True,
        )
        if torch.cuda.is_available():
            MODEL.to("cuda:0")
        PROCESSOR = AutoProcessor.from_pretrained(model_path)
        SERVER_CONFIG["pipeline"] = pipeline(
            "automatic-speech-recognition",
            model=MODEL,
            tokenizer=PROCESSOR.tokenizer,
            feature_extractor=PROCESSOR.feature_extractor,
            torch_dtype=dtype,
            device=0 if torch.cuda.is_available() else -1,
        )
        LOAD_ERROR = None
    except Exception as exc:  # pragma: no cover - depends on host GPU/runtime
        LOAD_ERROR = f"{type(exc).__name__}: {exc}"
        raise


def load_video_model() -> None:
    global LOAD_ERROR, MODEL, PROCESSOR
    if MODEL is not None and PROCESSOR is not None:
        return
    try:
        import torch
        from transformers import AutoProcessor

        try:
            from transformers import Qwen2_5_VLForConditionalGeneration as VisionModel
        except ImportError:
            from transformers import AutoModelForVision2Seq as VisionModel

        model_path = SERVER_CONFIG["model_path"]
        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
        MODEL = VisionModel.from_pretrained(
            model_path,
            torch_dtype=dtype,
            device_map="auto",
            trust_remote_code=True,
        )
        PROCESSOR = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
        LOAD_ERROR = None
    except Exception as exc:  # pragma: no cover - depends on host GPU/runtime
        LOAD_ERROR = f"{type(exc).__name__}: {exc}"
        raise


def load_image_model() -> None:
    global LOAD_ERROR, MODEL
    if MODEL is not None and MODEL_LOAD_STATE.get("image", {}).get("status") == "loaded":
        return
    with IMAGE_MODEL_LOCK:
        if MODEL is not None and MODEL_LOAD_STATE.get("image", {}).get("status") == "loaded":
            return
        if MODEL is not None:
            set_load_state("image", status="loaded", recovered_at=time.time(), error=None)
            return
        started = time.perf_counter()
        set_load_state("image", status="loading", started_at=time.time(), error=None)
        try:
            if not has_active_image_job("edit"):
                unload_image_edit_model()
            model_path = SERVER_CONFIG["model_path"]
            MODEL, torch, class_name = load_diffusers_pipeline(model_path, edit=False)
            set_load_state("image", status="moving_to_gpu" if torch.cuda.is_available() else "loading", error=None)
            prepare_diffusers_pipeline(MODEL, torch_module=torch, class_name=class_name)
            LOAD_ERROR = None
            set_load_state("image", status="loaded", pipeline_class=class_name, latency_ms=round((time.perf_counter() - started) * 1000), error=None)
        except Exception as exc:  # pragma: no cover - depends on host GPU/runtime
            MODEL = None
            LOAD_ERROR = f"{type(exc).__name__}: {exc}"
            set_load_state("image", status="error", error=LOAD_ERROR, latency_ms=round((time.perf_counter() - started) * 1000))
            raise


def load_image_edit_model() -> None:
    global LOAD_ERROR, EDIT_MODEL
    if EDIT_MODEL is not None and MODEL_LOAD_STATE.get("image_edit", {}).get("status") == "loaded":
        return
    with IMAGE_EDIT_MODEL_LOCK:
        if EDIT_MODEL is not None and MODEL_LOAD_STATE.get("image_edit", {}).get("status") == "loaded":
            return
        if EDIT_MODEL is not None:
            set_load_state("image_edit", status="loaded", recovered_at=time.time(), error=None)
            return
        started = time.perf_counter()
        set_load_state("image_edit", status="loading", started_at=time.time(), error=None)
        try:
            if not has_active_image_job("generate"):
                unload_image_model()
            model_path = SERVER_CONFIG.get("edit_model_path") or SERVER_CONFIG["model_path"]
            EDIT_MODEL, torch, class_name = load_diffusers_pipeline(model_path, edit=True)
            set_load_state("image_edit", status="moving_to_gpu" if torch.cuda.is_available() else "loading", error=None)
            prepare_diffusers_pipeline(EDIT_MODEL, torch_module=torch, class_name=class_name)
            LOAD_ERROR = None
            set_load_state("image_edit", status="loaded", pipeline_class=class_name, latency_ms=round((time.perf_counter() - started) * 1000), error=None)
        except Exception as exc:  # pragma: no cover - depends on host GPU/runtime
            EDIT_MODEL = None
            LOAD_ERROR = f"{type(exc).__name__}: {exc}"
            set_load_state("image_edit", status="error", error=LOAD_ERROR, latency_ms=round((time.perf_counter() - started) * 1000))
            raise


def generate_translation(messages: list[dict], max_tokens: int, temperature: float) -> str:
    import torch

    load_translation_model()
    assert MODEL is not None
    assert TOKENIZER is not None

    prompt = last_user_message(messages)
    if not prompt:
        prompt = "Translate the input text."
    system = (
        "You are a precise translation engine. Preserve formatting, numbers, JSON keys, "
        "and technical terms unless the user explicitly asks otherwise."
    )
    chat = [{"role": "system", "content": system}, {"role": "user", "content": prompt}]
    input_ids = TOKENIZER.apply_chat_template(chat, add_generation_prompt=True, return_tensors="pt").to(MODEL.device)
    outputs = MODEL.generate(
        input_ids,
        max_new_tokens=max_tokens,
        do_sample=temperature > 0,
        temperature=max(temperature, 0.01),
        pad_token_id=TOKENIZER.eos_token_id,
    )
    generated = outputs[0][input_ids.shape[-1] :]
    return TOKENIZER.decode(generated, skip_special_tokens=True).strip()


def decode_data_url(value: str, *, field_name: str) -> bytes:
    data = str(value or "")
    if "," in data and data.split(",", 1)[0].startswith("data:"):
        data = data.split(",", 1)[1]
    if not data:
        raise ValueError(f"{field_name} is required")
    return base64.b64decode(data)


def transcribe_audio(payload: dict) -> dict:
    if SERVER_CONFIG["plugin"] != "audio":
        raise ValueError("audio transcription is available only for the audio plugin")
    suffix = Path(str(payload.get("filename") or "audio.wav")).suffix or ".wav"
    load_audio_model()
    pipeline_obj = SERVER_CONFIG.get("pipeline")
    if pipeline_obj is None:
        raise RuntimeError("audio pipeline is not initialized")

    audio_bytes = decode_data_url(str(payload.get("audio_base64") or ""), field_name="audio_base64")
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=True) as handle:
        handle.write(audio_bytes)
        handle.flush()
        kwargs = {}
        language = payload.get("language")
        task = payload.get("task")
        if language or task:
            kwargs["generate_kwargs"] = {}
            if language:
                kwargs["generate_kwargs"]["language"] = str(language)
            if task:
                kwargs["generate_kwargs"]["task"] = str(task)
        result = pipeline_obj(handle.name, **kwargs)
    return {
        "text": str(result.get("text") or "").strip(),
        "model": SERVER_CONFIG["served_model_name"],
        "bytes": len(audio_bytes),
    }


def analyze_image(payload: dict) -> dict:
    if SERVER_CONFIG["plugin"] != "video":
        raise ValueError("vision analysis is available only for the video plugin")

    import torch
    from PIL import Image

    load_video_model()
    assert MODEL is not None
    assert PROCESSOR is not None

    prompt = str(payload.get("prompt") or "").strip() or "Опиши изображение и перечисли важные детали."
    max_tokens = int(payload.get("max_tokens") or 512)
    image_bytes = decode_data_url(str(payload.get("image_base64") or ""), field_name="image_base64")
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt},
            ],
        }
    ]

    text = PROCESSOR.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = PROCESSOR(text=[text], images=[image], return_tensors="pt")
    device = getattr(MODEL, "device", None)
    if device is not None:
        inputs = {key: value.to(device) if hasattr(value, "to") else value for key, value in inputs.items()}

    with torch.inference_mode():
        outputs = MODEL.generate(**inputs, max_new_tokens=max_tokens)
    input_length = inputs["input_ids"].shape[-1]
    generated = outputs[0][input_length:]
    answer = PROCESSOR.decode(generated, skip_special_tokens=True).strip()
    return {
        "text": answer,
        "model": SERVER_CONFIG["served_model_name"],
        "bytes": len(image_bytes),
        "width": image.width,
        "height": image.height,
    }


def encode_image_data_url(image) -> str:
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
    return f"data:image/png;base64,{encoded}"


def clamp_image_dimension(value: object, fallback: int = 1024) -> int:
    try:
        number = int(value)
    except (TypeError, ValueError):
        number = fallback
    number = max(512, min(1536, number))
    return number - (number % 8)


def optional_seed(value: object) -> int | None:
    if value in {None, ""}:
        return None
    return int(value)


def image_generator(seed: int | None):
    if seed is None:
        return None
    import torch

    device = "cuda" if torch.cuda.is_available() else "cpu"
    return torch.Generator(device=device).manual_seed(seed)


def gpu_runtime_status() -> dict:
    try:
        import torch

        if not torch.cuda.is_available():
            return {"available": False, "device_count": 0, "summary": "cuda unavailable"}
        devices = []
        for index in range(torch.cuda.device_count()):
            total = torch.cuda.get_device_properties(index).total_memory
            allocated = torch.cuda.memory_allocated(index)
            reserved = torch.cuda.memory_reserved(index)
            free = max(total - reserved, 0)
            devices.append(
                {
                    "index": index,
                    "name": torch.cuda.get_device_name(index),
                    "memory_total_mib": round(total / 1024 / 1024),
                    "memory_allocated_mib": round(allocated / 1024 / 1024),
                    "memory_reserved_mib": round(reserved / 1024 / 1024),
                    "memory_free_mib": round(free / 1024 / 1024),
                }
            )
        summary = " · ".join(
            f"GPU{item['index']} {item['memory_free_mib']}/{item['memory_total_mib']} MiB free"
            for item in devices
        )
        return {"available": True, "device_count": len(devices), "devices": devices, "summary": summary}
    except Exception as exc:  # pragma: no cover - depends on runtime
        return {"available": False, "device_count": 0, "summary": str(exc), "error": str(exc)}


def image_job_public(job: dict) -> dict:
    public = dict(job)
    payload = public.get("payload")
    if isinstance(payload, dict):
        public["payload"] = {
            key: value
            for key, value in payload.items()
            if key not in {"image_base64", "mask_base64"}
        }
    return public


def get_image_job(job_id: str) -> dict | None:
    with IMAGE_JOB_LOCK:
        job = IMAGE_JOBS.get(job_id)
        return image_job_public(job) if job else None


def get_image_job_internal(job_id: str) -> dict | None:
    with IMAGE_JOB_LOCK:
        job = IMAGE_JOBS.get(job_id)
        return dict(job) if job else None


def update_image_job(job_id: str, **fields: object) -> dict | None:
    with IMAGE_JOB_LOCK:
        job = IMAGE_JOBS.get(job_id)
        if not job:
            return None
        job.update(fields)
        job["updated_at"] = time.time()
        return image_job_public(job)


def list_image_jobs(limit: int = 20) -> list[dict]:
    with IMAGE_JOB_LOCK:
        jobs = sorted(IMAGE_JOBS.values(), key=lambda item: float(item.get("created_at") or 0), reverse=True)
        return [image_job_public(job) for job in jobs[:limit]]


def ensure_not_cancelled(cancel_event: threading.Event | None) -> None:
    if cancel_event is not None and cancel_event.is_set():
        raise RuntimeError("cancelled")


def diffusers_cancel_callback(cancel_event: threading.Event | None):
    if cancel_event is None:
        return None

    def callback(pipe, step, timestep, callback_kwargs):
        if cancel_event.is_set():
            setattr(pipe, "_interrupt", True)
        return callback_kwargs

    return callback


def call_diffusers_pipeline(pipe, *, cancel_event: threading.Event | None = None, **kwargs):
    callback = diffusers_cancel_callback(cancel_event)
    if callback is not None:
        kwargs["callback_on_step_end"] = callback
    try:
        return pipe(**kwargs)
    except TypeError:
        kwargs.pop("callback_on_step_end", None)
        ensure_not_cancelled(cancel_event)
        return pipe(**kwargs)


def generate_image(payload: dict, cancel_event: threading.Event | None = None) -> dict:
    if SERVER_CONFIG["plugin"] != "image":
        raise ValueError("image generation is available only for the image plugin")
    load_image_model()
    assert MODEL is not None
    ensure_not_cancelled(cancel_event)

    prompt = str(payload.get("prompt") or "").strip()
    if not prompt:
        raise ValueError("prompt is required")
    width = clamp_image_dimension(payload.get("width"), 1024)
    height = clamp_image_dimension(payload.get("height"), 1024)
    steps = max(1, min(80, int(payload.get("steps") or 28)))
    guidance_scale = float(payload.get("guidance_scale") or 6.0)
    seed = optional_seed(payload.get("seed"))
    ensure_not_cancelled(cancel_event)
    class_name = diffusers_class_name(SERVER_CONFIG["model_path"])
    kwargs = {
        "prompt": prompt,
        "negative_prompt": str(payload.get("negative_prompt") or "") or None,
        "width": width,
        "height": height,
        "num_inference_steps": steps,
        "generator": image_generator(seed),
    }
    if is_qwen_image_class(class_name):
        kwargs["true_cfg_scale"] = guidance_scale
    else:
        kwargs["guidance_scale"] = guidance_scale
    result = call_diffusers_pipeline(MODEL, cancel_event=cancel_event, **kwargs)
    ensure_not_cancelled(cancel_event)
    image = result.images[0]
    return {
        "image_base64": encode_image_data_url(image),
        "model": SERVER_CONFIG["served_model_name"],
        "width": image.width,
        "height": image.height,
        "seed": seed,
    }


def edit_image(payload: dict, cancel_event: threading.Event | None = None) -> dict:
    if SERVER_CONFIG["plugin"] != "image":
        raise ValueError("image editing is available only for the image plugin")

    from PIL import Image

    load_image_edit_model()
    assert EDIT_MODEL is not None
    ensure_not_cancelled(cancel_event)

    prompt = str(payload.get("prompt") or "").strip()
    if not prompt:
        raise ValueError("prompt is required")
    image_bytes = decode_data_url(str(payload.get("image_base64") or ""), field_name="image_base64")
    width = clamp_image_dimension(payload.get("width"), 1024)
    height = clamp_image_dimension(payload.get("height"), 1024)
    steps = max(1, min(80, int(payload.get("steps") or 28)))
    guidance_scale = float(payload.get("guidance_scale") or 6.0)
    strength = max(0.0, min(1.0, float(payload.get("strength") or 0.95)))
    seed = optional_seed(payload.get("seed"))
    source_image = Image.open(io.BytesIO(image_bytes)).convert("RGB").resize((width, height))
    ensure_not_cancelled(cancel_event)
    class_name = diffusers_class_name(SERVER_CONFIG.get("edit_model_path") or SERVER_CONFIG["model_path"])
    kwargs = {
        "prompt": prompt,
        "negative_prompt": str(payload.get("negative_prompt") or "") or None,
        "image": source_image,
        "width": width,
        "height": height,
        "num_inference_steps": steps,
        "generator": image_generator(seed),
    }
    if is_qwen_image_class(class_name):
        kwargs["true_cfg_scale"] = guidance_scale
    else:
        mask_bytes = decode_data_url(str(payload.get("mask_base64") or ""), field_name="mask_base64")
        kwargs["mask_image"] = Image.open(io.BytesIO(mask_bytes)).convert("L").resize((width, height))
        kwargs["guidance_scale"] = guidance_scale
        kwargs["strength"] = strength
    result = call_diffusers_pipeline(EDIT_MODEL, cancel_event=cancel_event, **kwargs)
    ensure_not_cancelled(cancel_event)
    image = result.images[0]
    return {
        "image_base64": encode_image_data_url(image),
        "model": SERVER_CONFIG["served_model_name"],
        "width": image.width,
        "height": image.height,
        "seed": seed,
    }


def run_image_job(job_id: str) -> None:
    job = get_image_job_internal(job_id)
    if not job:
        return
    cancel_event = IMAGE_JOB_CANCEL_EVENTS[job_id]
    operation = str(job.get("operation") or "generate")
    payload = dict(job.get("payload") or {})
    started = time.perf_counter()
    load_key = "image_edit" if operation == "edit" else "image"
    model_loaded = MODEL_LOAD_STATE.get(load_key, {}).get("status") == "loaded"
    update_image_job(
        job_id,
        status="running" if model_loaded else "loading_model",
        started_at=time.time(),
        load_state=MODEL_LOAD_STATE.get(load_key),
    )
    try:
        ensure_not_cancelled(cancel_event)
        if operation == "edit":
            load_image_edit_model()
            update_image_job(job_id, status="running", load_state=MODEL_LOAD_STATE.get(load_key))
            result = edit_image(payload, cancel_event=cancel_event)
        else:
            load_image_model()
            update_image_job(job_id, status="running", load_state=MODEL_LOAD_STATE.get(load_key))
            result = generate_image(payload, cancel_event=cancel_event)
        latency_ms = round((time.perf_counter() - started) * 1000)
        if cancel_event.is_set():
            update_image_job(job_id, status="cancelled", latency_ms=latency_ms)
        else:
            result["latency_ms"] = latency_ms
            update_image_job(job_id, status="completed", result=result, latency_ms=latency_ms)
    except RuntimeError as exc:
        latency_ms = round((time.perf_counter() - started) * 1000)
        if str(exc).lower() == "cancelled":
            update_image_job(job_id, status="cancelled", latency_ms=latency_ms)
        else:
            update_image_job(job_id, status="error", error=str(exc), error_type=type(exc).__name__, latency_ms=latency_ms)
    except Exception as exc:  # pragma: no cover - depends on host GPU/runtime
        latency_ms = round((time.perf_counter() - started) * 1000)
        update_image_job(job_id, status="error", error=str(exc), error_type=type(exc).__name__, latency_ms=latency_ms)


class PluginHandler(BaseHTTPRequestHandler):
    server_version = "LLMTransformersPlugin/1.0"

    def log_message(self, fmt: str, *args: object) -> None:
        print(f"{self.address_string()} - {fmt % args}", file=sys.stderr)

    def do_GET(self) -> None:
        if self.path == "/health":
            json_response(
                self,
                200 if LOAD_ERROR is None else 503,
                {
                    "status": "ok" if LOAD_ERROR is None else "error",
                    "plugin": SERVER_CONFIG["plugin"],
                    "model": SERVER_CONFIG["served_model_name"],
                    "model_path": SERVER_CONFIG["model_path"],
                    "edit_model_path": SERVER_CONFIG.get("edit_model_path"),
                    "loaded": MODEL is not None,
                    "edit_loaded": EDIT_MODEL is not None,
                    "load_state": MODEL_LOAD_STATE,
                    "gpu": gpu_runtime_status(),
                    "error": LOAD_ERROR,
                },
            )
            return
        if self.path == "/v1/images/jobs":
            json_response(self, 200, {"jobs": list_image_jobs()})
            return
        if self.path.startswith("/v1/images/jobs/"):
            job_id = self.path.rsplit("/", 1)[-1].split("?", 1)[0]
            job = get_image_job(job_id)
            if not job:
                json_response(self, 404, {"error": {"message": "job not found", "type": "not_found"}})
                return
            json_response(self, 200, job)
            return
        if self.path == "/v1/models":
            json_response(
                self,
                200,
                {
                    "object": "list",
                    "data": [
                        {
                            "id": SERVER_CONFIG["served_model_name"],
                            "object": "model",
                            "created": 0,
                            "owned_by": "local",
                        }
                    ],
                },
            )
            return
        json_response(self, 404, {"error": "not found"})

    def do_POST(self) -> None:
        if self.path == "/v1/images/jobs":
            try:
                body = read_body(self)
                operation = str(body.get("operation") or "generate")
                if operation not in {"generate", "edit"}:
                    raise ValueError("operation must be generate or edit")
                payload = body.get("payload") if isinstance(body.get("payload"), dict) else body
                job_id = str(body.get("job_id") or uuid.uuid4())
                cancel_event = threading.Event()
                with IMAGE_JOB_LOCK:
                    IMAGE_JOBS[job_id] = {
                        "id": job_id,
                        "operation": operation,
                        "status": "queued",
                        "created_at": time.time(),
                        "updated_at": time.time(),
                        "payload": dict(payload),
                    }
                    IMAGE_JOB_CANCEL_EVENTS[job_id] = cancel_event
                thread = threading.Thread(target=run_image_job, args=(job_id,), daemon=True)
                thread.start()
                json_response(self, 202, get_image_job(job_id) or {"id": job_id, "status": "queued"})
            except Exception as exc:
                json_response(self, 400, {"error": {"message": str(exc), "type": type(exc).__name__}})
            return

        if self.path.startswith("/v1/images/jobs/") and self.path.endswith("/cancel"):
            job_id = self.path.removeprefix("/v1/images/jobs/").removesuffix("/cancel").strip("/")
            event = IMAGE_JOB_CANCEL_EVENTS.get(job_id)
            if event is None:
                json_response(self, 404, {"error": {"message": "job not found", "type": "not_found"}})
                return
            event.set()
            job = get_image_job(job_id)
            if job and job.get("status") in {"queued", "loading_model", "running"}:
                update_image_job(job_id, status="cancel_requested")
            json_response(self, 200, get_image_job(job_id) or {"id": job_id, "status": "cancel_requested"})
            return

        if self.path == "/v1/images/generations":
            started = time.perf_counter()
            try:
                result = generate_image(read_body(self))
            except Exception as exc:  # pragma: no cover - depends on host GPU/runtime
                json_response(self, 503 if LOAD_ERROR else 400, {"error": {"message": str(exc), "type": type(exc).__name__}})
                return
            result["latency_ms"] = round((time.perf_counter() - started) * 1000)
            json_response(self, 200, result)
            return

        if self.path == "/v1/images/edits":
            started = time.perf_counter()
            try:
                result = edit_image(read_body(self))
            except Exception as exc:  # pragma: no cover - depends on host GPU/runtime
                json_response(self, 503 if LOAD_ERROR else 400, {"error": {"message": str(exc), "type": type(exc).__name__}})
                return
            result["latency_ms"] = round((time.perf_counter() - started) * 1000)
            json_response(self, 200, result)
            return

        if self.path == "/v1/vision/analyze":
            started = time.perf_counter()
            try:
                result = analyze_image(read_body(self))
            except Exception as exc:  # pragma: no cover - depends on host GPU/runtime
                json_response(self, 503 if LOAD_ERROR else 400, {"error": {"message": str(exc), "type": type(exc).__name__}})
                return
            result["latency_ms"] = round((time.perf_counter() - started) * 1000)
            json_response(self, 200, result)
            return

        if self.path == "/v1/audio/transcriptions":
            started = time.perf_counter()
            try:
                result = transcribe_audio(read_body(self))
            except Exception as exc:  # pragma: no cover - depends on host GPU/runtime
                json_response(self, 503 if LOAD_ERROR else 400, {"error": {"message": str(exc), "type": type(exc).__name__}})
                return
            result["latency_ms"] = round((time.perf_counter() - started) * 1000)
            json_response(self, 200, result)
            return

        if self.path != "/v1/chat/completions":
            json_response(self, 404, {"error": "not found"})
            return
        body = read_body(self)
        messages = body.get("messages") or []
        max_tokens = int(body.get("max_tokens") or 512)
        temperature = float(body.get("temperature") or 0.0)

        if SERVER_CONFIG["plugin"] != "translation":
            json_response(
                self,
                501,
                {
                    "error": {
                        "message": f"{SERVER_CONFIG['plugin']} chat generation is not implemented in this service yet",
                        "type": "not_implemented",
                    }
                },
            )
            return

        try:
            content = generate_translation(messages, max_tokens=max_tokens, temperature=temperature)
        except Exception as exc:  # pragma: no cover - depends on host GPU/runtime
            json_response(self, 503, {"error": {"message": str(exc), "type": type(exc).__name__}})
            return

        json_response(
            self,
            200,
            {
                "id": f"chatcmpl-{uuid.uuid4().hex}",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": SERVER_CONFIG["served_model_name"],
                "choices": [
                    {
                        "index": 0,
                        "message": {"role": "assistant", "content": content},
                        "finish_reason": "stop",
                    }
                ],
            },
        )


def main() -> int:
    parser = argparse.ArgumentParser(description="Small OpenAI-like service for local Transformers plugin models.")
    parser.add_argument("--plugin", required=True, choices=["translation", "audio", "video", "image"])
    parser.add_argument("--model-path", default=os.environ.get("MODEL_PATH"))
    parser.add_argument("--edit-model-path", default=os.environ.get("EDIT_MODEL_PATH"))
    parser.add_argument("--served-model-name", default=os.environ.get("SERVED_MODEL_NAME"))
    parser.add_argument("--host", default=os.environ.get("HOST", "0.0.0.0"))
    parser.add_argument("--port", type=int, default=int(os.environ.get("PORT", "8010")))
    parser.add_argument("--load-on-start", action="store_true", default=os.environ.get("LOAD_ON_START") == "1")
    parser.add_argument("--background-load-on-start", action="store_true", default=os.environ.get("BACKGROUND_LOAD_ON_START") == "1")
    args = parser.parse_args()

    if not args.model_path:
        parser.error("--model-path or MODEL_PATH is required")
    if not args.served_model_name:
        args.served_model_name = Path(args.model_path).name

    SERVER_CONFIG.update(
        {
            "plugin": args.plugin,
            "model_path": args.model_path,
            "edit_model_path": args.edit_model_path,
            "served_model_name": args.served_model_name,
        }
    )

    def preload() -> None:
        try:
            if args.plugin == "translation":
                load_translation_model()
            elif args.plugin == "audio":
                load_audio_model()
            elif args.plugin == "video":
                load_video_model()
            elif args.plugin == "image":
                load_image_model()
        except Exception as exc:  # pragma: no cover - depends on host GPU/runtime
            print(f"preload failed: {type(exc).__name__}: {exc}", file=sys.stderr, flush=True)

    if args.load_on_start and not args.background_load_on_start:
        preload()

    server = ThreadingHTTPServer((args.host, args.port), PluginHandler)
    print(f"{args.plugin} service listening on {args.host}:{args.port} model={args.served_model_name}", flush=True)
    if args.load_on_start and args.background_load_on_start:
        threading.Thread(target=preload, daemon=True).start()
    server.serve_forever()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
