#!/usr/bin/env python3
"""
DaSiWa I2V/FLF2V API Server для ComfyUI.
Работает рядом с ComfyUI на той же машине.

Принимает HTTP запросы с HMAC авторизацией,
отправляет workflow в ComfyUI, возвращает видео.
"""

import os
import sys
import json
import uuid
import time
import shutil
import base64
import random
import logging
import binascii
import subprocess
import urllib.request
import urllib.parse
import websocket as ws_client

from flask import Flask, request, jsonify
from hmac_auth import verify_request

# ============================================================================
# Конфигурация
# ============================================================================

COMFY_HOST = os.getenv("COMFY_HOST", "127.0.0.1")
COMFY_PORT = os.getenv("COMFY_PORT", "8188")
API_PORT = int(os.getenv("API_PORT", "8080"))

WORKFLOW_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "workflow_api.json")
KEYS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "keys.json")

COMFY_OUTPUT_DIR = os.getenv("COMFY_OUTPUT_DIR", "/ComfyUI/output")

# ============================================================================
# Инициализация
# ============================================================================

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s"
)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# Загрузка ключей
if not os.path.exists(KEYS_FILE):
    logger.error(f"❌ Файл ключей не найден: {KEYS_FILE}")
    logger.error("   Запусти: python generate_keys.py")
    sys.exit(1)

with open(KEYS_FILE, "r") as f:
    keys = json.load(f)
    CLIENT_ID = keys["client_id"]
    SECRET_KEY = keys["secret_key"]

logger.info(f"🔐 Ключи загружены. Client ID: {CLIENT_ID}")

# Множество использованных nonce (защита от replay-атак)
used_nonces = set()

# WebSocket client ID
ws_client_id = str(uuid.uuid4())


# ============================================================================
# Утилиты
# ============================================================================

def to_nearest_multiple_of_16(value):
    """Округляет значение до ближайшего кратного 16."""
    try:
        numeric_value = float(value)
    except Exception:
        raise ValueError(f"width/height value is not a number: {value}")
    adjusted = int(round(numeric_value / 16.0) * 16)
    return max(adjusted, 16)


def save_base64_to_file(base64_data, temp_dir, filename):
    """Сохраняет base64 данные в файл."""
    decoded = base64.b64decode(base64_data)
    os.makedirs(temp_dir, exist_ok=True)
    file_path = os.path.abspath(os.path.join(temp_dir, filename))
    with open(file_path, "wb") as f:
        f.write(decoded)
    return file_path


def download_file(url, temp_dir, filename):
    """Скачивает файл по URL."""
    os.makedirs(temp_dir, exist_ok=True)
    file_path = os.path.abspath(os.path.join(temp_dir, filename))
    result = subprocess.run(
        ["wget", "-O", file_path, "--no-verbose", url],
        capture_output=True, text=True, timeout=120
    )
    if result.returncode != 0:
        raise RuntimeError(f"Download failed: {result.stderr}")
    return file_path


def process_image_input(job_input, prefix, temp_dir):
    """
    Обрабатывает входные данные изображения.
    prefix: "image" или "last_image"
    Возвращает (file_path, True) или (None, False)
    """
    path_key = f"{prefix}_path"
    url_key = f"{prefix}_url"
    b64_key = f"{prefix}_base64"

    if path_key in job_input and job_input[path_key]:
        return job_input[path_key], True
    elif url_key in job_input and job_input[url_key]:
        return download_file(job_input[url_key], temp_dir, f"{prefix}.png"), True
    elif b64_key in job_input and job_input[b64_key]:
        return save_base64_to_file(job_input[b64_key], temp_dir, f"{prefix}.png"), True

    return None, False


def queue_prompt(prompt):
    """Отправляет prompt в ComfyUI."""
    url = f"http://{COMFY_HOST}:{COMFY_PORT}/prompt"
    data = json.dumps({"prompt": prompt, "client_id": ws_client_id}).encode("utf-8")
    req = urllib.request.Request(url, data=data)
    return json.loads(urllib.request.urlopen(req).read())


def get_history(prompt_id):
    """Получает историю выполнения prompt."""
    url = f"http://{COMFY_HOST}:{COMFY_PORT}/history/{prompt_id}"
    with urllib.request.urlopen(url) as response:
        return json.loads(response.read())


def generate_video(prompt):
    """Подключается к ComfyUI по WebSocket, запускает генерацию, ждёт результат."""
    ws_url = f"ws://{COMFY_HOST}:{COMFY_PORT}/ws?clientId={ws_client_id}"

    ws = ws_client.WebSocket()
    ws.connect(ws_url)
    logger.info("🔌 WebSocket подключён к ComfyUI")

    prompt_id = queue_prompt(prompt)["prompt_id"]
    logger.info(f"📤 Prompt отправлен: {prompt_id}")

    # Ждём завершения
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            if message["type"] == "executing":
                data = message["data"]
                if data["node"] is None and data["prompt_id"] == prompt_id:
                    break

    ws.close()
    logger.info("✅ Генерация завершена")

    # Извлекаем видео
    history = get_history(prompt_id)[prompt_id]
    for node_id in history["outputs"]:
        node_output = history["outputs"][node_id]
        if "gifs" in node_output:
            for video in node_output["gifs"]:
                video_path = video["fullpath"]
                with open(video_path, "rb") as f:
                    video_b64 = base64.b64encode(f.read()).decode("utf-8")
                # Очистка
                try:
                    os.remove(video_path)
                except OSError:
                    pass
                return video_b64

    return None


# ============================================================================
# API Endpoints
# ============================================================================

@app.before_request
def check_hmac_auth():
    """Проверяет HMAC подпись для всех запросов кроме health check."""
    if request.path == "/health":
        return None

    body = request.get_data()
    headers = {
        "X-Client-Id": request.headers.get("X-Client-Id", ""),
        "X-Timestamp": request.headers.get("X-Timestamp", ""),
        "X-Nonce": request.headers.get("X-Nonce", ""),
        "X-Signature": request.headers.get("X-Signature", ""),
    }

    is_valid, error = verify_request(body, headers, SECRET_KEY, CLIENT_ID, used_nonces)
    if not is_valid:
        logger.warning(f"🚫 Auth failed: {error} from {request.remote_addr}")
        return jsonify({"error": "Unauthorized", "detail": error}), 401


@app.route("/health", methods=["GET"])
def health():
    """Health check — без авторизации."""
    try:
        url = f"http://{COMFY_HOST}:{COMFY_PORT}/"
        urllib.request.urlopen(url, timeout=5)
        comfy_status = "ok"
    except Exception:
        comfy_status = "unavailable"

    return jsonify({
        "status": "ok",
        "comfyui": comfy_status,
        "timestamp": int(time.time())
    })


@app.route("/generate", methods=["POST"])
def generate():
    """Основной endpoint для генерации видео."""
    start_time = time.time()

    job_input = request.json or {}

    logger.info("=" * 60)
    logger.info("🎬 Новый запрос на генерацию")
    logger.info("=" * 60)

    # Логирование (без base64 данных)
    log_input = {k: v for k, v in job_input.items()
                 if not k.endswith("_base64")}
    logger.info(f"Параметры: {json.dumps(log_input, ensure_ascii=False)}")

    task_id = f"task_{uuid.uuid4().hex[:8]}"
    temp_dir = os.path.join("/tmp", task_id)

    try:
        # === Обработка изображений ===
        image_path, has_image = process_image_input(job_input, "image", temp_dir)
        if not has_image:
            return jsonify({"error": "No input image provided. Use image_base64, image_url, or image_path"}), 400

        last_image_path, use_flf2v = process_image_input(job_input, "last_image", temp_dir)

        mode = "FLF2V" if use_flf2v else "I2V"
        logger.info(f"🎬 Режим: {mode}")

        # === Загрузка workflow ===
        if not os.path.exists(WORKFLOW_FILE):
            return jsonify({"error": f"Workflow file not found: {WORKFLOW_FILE}"}), 500

        with open(WORKFLOW_FILE, "r") as f:
            prompt = json.load(f)

        # === Параметры генерации ===
        width = to_nearest_multiple_of_16(job_input.get("width", 528))
        height = to_nearest_multiple_of_16(job_input.get("height", 768))
        length = job_input.get("length", 81)
        steps = job_input.get("steps", 4)
        cfg = job_input.get("cfg", 1.0)
        seed = job_input.get("seed", -1)
        fps = job_input.get("fps", 16)
        sampler_name = job_input.get("sampler_name", "euler")
        scheduler = job_input.get("scheduler", "linear_quadratic")

        if seed == -1:
            seed = random.randint(0, 2**63 - 1)

        logger.info(f"📐 {width}x{height}, {length} frames, {steps} steps, CFG {cfg}, seed {seed}")

        # === Заполнение workflow ===

        # Positive prompt
        prompt["5"]["inputs"]["text"] = job_input.get("prompt", "")

        # Negative prompt
        negative_prompt = job_input.get("negative_prompt", prompt["6"]["inputs"]["text"])
        prompt["6"]["inputs"]["text"] = negative_prompt

        # First frame image
        prompt["7"]["inputs"]["image"] = image_path

        # FLF2V / I2V mode
        if use_flf2v and last_image_path:
            prompt["15"]["inputs"]["image"] = last_image_path
            logger.info(f"🎬 FLF2V: last frame = {last_image_path}")
        else:
            prompt["8"]["class_type"] = "WanImageToVideo"
            if "end_image" in prompt["8"]["inputs"]:
                del prompt["8"]["inputs"]["end_image"]
            if "15" in prompt:
                del prompt["15"]
            logger.info("🎬 I2V: single image mode")

        # Video dimensions
        prompt["8"]["inputs"]["width"] = width
        prompt["8"]["inputs"]["height"] = height
        prompt["8"]["inputs"]["length"] = length

        # KSampler High
        prompt["11"]["inputs"]["noise_seed"] = seed
        prompt["11"]["inputs"]["steps"] = steps
        prompt["11"]["inputs"]["cfg"] = cfg
        prompt["11"]["inputs"]["sampler_name"] = sampler_name
        prompt["11"]["inputs"]["scheduler"] = scheduler
        prompt["11"]["inputs"]["end_at_step"] = steps // 2

        # KSampler Low
        prompt["12"]["inputs"]["noise_seed"] = seed
        prompt["12"]["inputs"]["steps"] = steps
        prompt["12"]["inputs"]["cfg"] = cfg
        prompt["12"]["inputs"]["sampler_name"] = sampler_name
        prompt["12"]["inputs"]["scheduler"] = scheduler
        prompt["12"]["inputs"]["start_at_step"] = steps // 2

        # Video output
        prompt["14"]["inputs"]["frame_rate"] = fps

        # === Генерация ===
        video_b64 = generate_video(prompt)

        if not video_b64:
            return jsonify({"error": "Video generation failed — no output"}), 500

        elapsed = time.time() - start_time
        logger.info(f"✅ Видео сгенерировано за {elapsed:.1f}s")

        return jsonify({
            "video": video_b64,
            "seed": seed,
            "mode": mode,
            "elapsed": round(elapsed, 1)
        })

    except Exception as e:
        logger.error(f"❌ Ошибка: {e}", exc_info=True)
        return jsonify({"error": str(e)}), 500

    finally:
        # Очистка temp файлов
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir, ignore_errors=True)
        # Очистка output ComfyUI
        try:
            if os.path.exists(COMFY_OUTPUT_DIR):
                for fname in os.listdir(COMFY_OUTPUT_DIR):
                    fpath = os.path.join(COMFY_OUTPUT_DIR, fname)
                    if os.path.isfile(fpath):
                        os.unlink(fpath)
                    elif os.path.isdir(fpath):
                        shutil.rmtree(fpath)
        except Exception:
            pass


# ============================================================================
# Запуск
# ============================================================================

if __name__ == "__main__":
    logger.info("=" * 60)
    logger.info("🚀 DaSiWa API Server")
    logger.info(f"   ComfyUI:  http://{COMFY_HOST}:{COMFY_PORT}")
    logger.info(f"   API Port: {API_PORT}")
    logger.info(f"   Workflow: {WORKFLOW_FILE}")
    logger.info("=" * 60)

    # Проверяем подключение к ComfyUI
    try:
        urllib.request.urlopen(f"http://{COMFY_HOST}:{COMFY_PORT}/", timeout=5)
        logger.info("✅ ComfyUI доступен")
    except Exception:
        logger.warning("⚠️  ComfyUI недоступен — запросы будут ждать")

    app.run(host="0.0.0.0", port=API_PORT, debug=False)