import asyncio
import json
import os
import time
import urllib.request
import urllib.parse
import folder_paths
from aiohttp import web
import server
import re

from . import utils


prompt_server = server.PromptServer.instance
main_loop = asyncio.get_event_loop()

def sanitize_filename(filename):
    filename = filename.replace("..", "").replace("\0", "")
    illegal_chars = r'<>:"/\\|?*\t\n\r'
    sanitized = re.sub(f"[{re.escape(illegal_chars)}]", "_", filename)
    name, ext = os.path.splitext(sanitized)
    name = name[:150]
    return f"{name}{ext}"


@prompt_server.routes.get("/civitai_utils/get_db_stats")
async def get_db_stats(request):
    try:
        stats = utils.db_manager.get_db_stats()
        return web.json_response({"status": "ok", "stats": stats})
    except Exception as e:
        print(f"[Civitai Utils] Error getting DB stats: {e}")
        return web.json_response({"status": "error", "message": str(e)}, status=500)


@prompt_server.routes.get("/civitai_utils/get_scanned_models")
async def get_scanned_models(request):
    """获取数据库中已扫描模型列表的API"""
    model_type = request.query.get("model_type")
    if not model_type or model_type not in ["checkpoints", "loras"]:
        return web.json_response(
            {"status": "error", "message": "Invalid model_type"}, status=400
        )
    try:
        # force_sync=False 避免不必要的重复扫描
        model_list = utils.get_model_filenames_from_db(model_type, force_sync=False)
        return web.json_response({"status": "ok", "models": model_list})
    except Exception as e:
        print(f"[Civitai Utils] Error getting scanned models: {e}")
        return web.json_response({"status": "error", "message": str(e)}, status=500)


@prompt_server.routes.get("/civitai_utils/check_legacy_cache")
async def check_legacy_cache(request):
    """检查旧版缓存文件是否存在的API"""
    try:
        exists = utils.check_legacy_cache_exists()
        return web.json_response({"exists": exists})
    except Exception as e:
        return web.json_response({"exists": False, "error": str(e)})


@prompt_server.routes.post("/civitai_utils/force_rescan")
async def force_rescan(request):
    try:
        data = await request.json()
        model_type = data.get("model_type")
        rehash_all = data.get("rehash_all", False)

        if not model_type or model_type not in ["checkpoints", "loras"]:
            return web.json_response(
                {"status": "error", "message": "Invalid model_type"}, status=400
            )

        # 如果是 rehash_all，我们需要先清空数据库中的 mtime
        if rehash_all:
            with utils.db_manager.get_connection() as conn:
                conn.execute(
                    "UPDATE versions SET local_mtime = 0 WHERE model_type = ?",
                    (model_type,),
                )

        # 将计时器清零以确保扫描执行
        utils.db_manager.set_setting(f"last_sync_{model_type}", 0)
        # 执行扫描并捕获结果
        scan_results = utils.sync_local_files_with_db(model_type, force=True)
        found_count = scan_results.get("found", 0)
        hashed_count = scan_results.get("hashed", 0)

        if rehash_all:
            message = (
                f"Re-hash complete! {hashed_count} {model_type} files were re-hashed."
            )
        elif found_count > 0:
            message = f"Rescan complete! Found and processed {hashed_count} new/modified {model_type} files."
        else:
            message = f"Rescan complete. No new or modified {model_type} files found."

        return web.json_response({"status": "ok", "message": message})
    except Exception as e:
        print(f"[Civitai Utils] Error forcing rescan: {e}")
        import traceback

        traceback.print_exc()
        return web.json_response({"status": "error", "message": str(e)}, status=500)


# 从旧版JSON迁移哈希的API
@prompt_server.routes.post("/civitai_utils/migrate_hashes")
async def migrate_hashes(request):
    try:
        results = utils.migrate_legacy_caches()
        return web.json_response({"status": "ok", "message": results["message"]})
    except Exception as e:
        print(f"[Civitai Utils] Error migrating hashes: {e}")
        import traceback

        traceback.print_exc()
        return web.json_response({"status": "error", "message": str(e)}, status=500)


@prompt_server.routes.post("/civitai_utils/clear_cache")
async def clear_cache(request):
    try:
        data = await request.json()
        cache_type = data.get("cache_type")

        if cache_type == "analysis":
            utils.db_manager.clear_analysis_cache()
            message = "Analyzer cache cleared successfully."
        elif cache_type == "api_responses":
            utils.db_manager.clear_api_responses()
            message = "API response cache cleared successfully."
        elif cache_type == "triggers":
            utils.db_manager.clear_all_triggers()
            message = "Trigger word cache cleared successfully."
        elif cache_type == "all":
            utils.db_manager.clear_analysis_cache()
            utils.db_manager.clear_api_responses()
            utils.db_manager.clear_all_triggers()
            message = "All caches have been cleared."
        else:
            return web.json_response(
                {"status": "error", "message": "Invalid cache type"}, status=400
            )

        return web.json_response({"status": "ok", "message": message})
    except Exception as e:
        print(f"[Civitai Utils] Error clearing cache: {e}")
        import traceback

        traceback.print_exc()
        return web.json_response({"status": "error", "message": str(e)}, status=500)


@prompt_server.routes.get("/civitai_recipe_finder/fetch_data")
async def fetch_data(request):
    try:

        model_type, model_filename, sort, nsfw_level, limit, filter_type = (
            request.query.get("model_type"),
            request.query.get("model_filename"),
            request.query.get("sort"),
            request.query.get("nsfw_level"),
            int(request.query.get("limit", 32)),
            request.query.get("filter_type"),
        )

        _, filename_to_hash = utils.get_local_model_maps(model_type)
        model_hash = filename_to_hash.get(model_filename)
        if not model_hash:
            raise FileNotFoundError(f"Model hash not found for: {model_filename} in type {model_type}")
        gallery_data = utils.fetch_civitai_data_by_hash(
            model_hash,
            sort,
            limit,
            nsfw_level,
            filter_type if filter_type != "all" else None,
        )
        return web.json_response({"status": "ok", "images": gallery_data})
    except Exception as e:
        print(f"Error in fetch_data: {e}")
        return web.json_response({"status": "error", "message": str(e)}, status=500)


@prompt_server.routes.post("/civitai_recipe_finder/set_selection")
async def set_selection(request):
    try:
        data = await request.json()
        node_id, item, download_image = (
            str(data.get("node_id")),
            data.get("item"),
            data.get("download_image", False),
        )
        selections = utils.load_selections()
        selections[node_id] = {"item": item, "download_image": download_image}
        utils.save_selections(selections)
        utils.db_manager.set_setting("last_selection_time", time.time())
        return web.json_response({"status": "ok"})
    except Exception as e:
        return web.json_response({"status": "error", "message": str(e)}, status=500)


@prompt_server.routes.post("/civitai_recipe_finder/save_original_image")
async def save_original_image(request):
    try:
        data = await request.json()
        image_url = data.get("url")
        if not image_url:
            return web.json_response(
                {"status": "error", "message": "URL is missing"}, status=400
            )
        clean_url = re.sub(r"/(width|height|fit|quality|format)=\w+", "", image_url)
        image_record = utils.db_manager.get_image_by_url(clean_url)
        if image_record and image_record["local_filename"]:
            local_path = os.path.join(
                folder_paths.get_output_directory(), image_record["local_filename"]
            )
            if os.path.exists(local_path):
                return web.json_response(
                    {
                        "status": "exists",
                        "message": f"Image already exists: {image_record['local_filename']}",
                    }
                )
        headers = {"User-Agent": "Mozilla/5.0"}
        req = urllib.request.Request(clean_url, headers=headers)

        base_filename = os.path.basename(urllib.parse.urlparse(clean_url).path)
        sanitized_base = sanitize_filename(base_filename)
        filename = f"civitai_{int(time.time())}_{sanitized_base}"

        output_path = os.path.join(folder_paths.get_output_directory(), filename)
        with (
            urllib.request.urlopen(req, timeout=20) as response,
            open(output_path, "wb") as f,
        ):
            f.write(response.read())
        utils.db_manager.add_downloaded_image(url=clean_url, local_filename=filename)
        return web.json_response(
            {"status": "ok", "message": f"Image saved as {filename}"}
        )
    except Exception as e:
        print(f"[Civitai Utils] Error in save_original_image: {e}")
        import traceback

        traceback.print_exc()
        return web.json_response({"status": "error", "message": str(e)}, status=500)


@prompt_server.routes.post("/civitai_recipe_finder/get_workflow_source")
async def get_workflow_source(request):
    try:
        data = await request.json()
        image_url = data.get("url")
        if not image_url:
            return web.Response(status=400, text="URL is missing")
        clean_url = re.sub(r"/(width|height|fit|quality|format)=\w+", "", image_url)
        image_record = utils.db_manager.get_image_by_url(clean_url)
        if image_record and image_record["local_filename"]:
            local_path = os.path.join(
                folder_paths.get_output_directory(), image_record["local_filename"]
            )
            if os.path.exists(local_path):
                with open(local_path, "rb") as f:
                    image_data = f.read()
                ext = os.path.splitext(image_record["local_filename"])[1].lower()
                content_type = {
                    ".png": "image/png",
                    ".jpg": "image/jpeg",
                    ".jpeg": "image/jpeg",
                    ".webp": "image/webp",
                }.get(ext, "image/png")
                return web.Response(body=image_data, content_type=content_type)
        headers = {"User-Agent": "Mozilla/5.0"}
        req = urllib.request.Request(clean_url, headers=headers)
        with urllib.request.urlopen(req, timeout=20) as response:
            image_data = response.read()
            content_type = response.headers.get("Content-Type", "image/png")
        filename_base = os.path.basename(urllib.parse.urlparse(clean_url).path)
        sanitized_base = sanitize_filename(filename_base)
        filename = f"civitai_{int(time.time())}_{sanitized_base}"
        if not any(
            filename.lower().endswith(ext) for ext in [".png", ".jpg", ".jpeg", ".webp"]
        ):
            ext = "." + content_type.split("/")[-1] if "/" in content_type else ".png"
            filename += ext
        output_path = os.path.join(folder_paths.get_output_directory(), filename)
        with open(output_path, "wb") as f:
            f.write(image_data)
        utils.db_manager.add_downloaded_image(url=clean_url, local_filename=filename)
        return web.Response(body=image_data, content_type=content_type)
    except Exception as e:
        return web.Response(status=500, text=str(e))


@prompt_server.routes.get("/civitai_utils/get_config")
async def get_config(request):
    api_key = utils.db_manager.get_setting("civitai_api_key")
    config = {
        "network_choice": utils.db_manager.get_setting("network_choice", "com"),
        "api_key_exists": bool(api_key)
    }
    return web.json_response(config)


@prompt_server.routes.post("/civitai_utils/set_config")
async def set_config(request):
    try:
        data = await request.json()
        if "network_choice" in data:
            utils.db_manager.set_setting("network_choice", data["network_choice"])
        if "api_key" in data:
            utils.db_manager.set_setting("civitai_api_key", data["api_key"])

        return web.json_response({"status": "ok"})
    except Exception as e:
        return web.json_response({"status": "error", "message": str(e)}, status=500)


@prompt_server.routes.get("/civitai_utils/get_local_models")
async def get_local_models(request):
    """
    一个强大、统一的API，一次性获取所有处理好的模型数据。
    """
    try:
        loop = asyncio.get_event_loop()
        force_refresh = request.query.get("force_refresh", "false").lower() == "true"

        # 将所有耗时的操作放入线程池，避免阻塞主服务器线程
        models = await loop.run_in_executor(
            None, utils.get_all_local_models_with_details, force_refresh
        )

        return web.json_response({"status": "ok", "models": models})
    except Exception as e:
        print(f"[Civitai Utils] FATAL ERROR in get_local_models API: {e}")
        import traceback

        traceback.print_exc()
        return web.json_response({"status": "error", "message": str(e)}, status=500)


@prompt_server.routes.get("/civitai_recipe_finder/get_local_hashes")
async def get_local_hashes(request):
    """
    一个轻量级的API，仅返回数据库中所有本地文件的哈希列表。
    """
    try:
        with utils.db_manager.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(
                "SELECT hash FROM versions WHERE hash IS NOT NULL AND local_path IS NOT NULL"
            )
            # 使用集合推导式以获得最佳性能
            hashes = {row["hash"] for row in cursor.fetchall()}
        return web.json_response({"status": "ok", "hashes": list(hashes)})
    except Exception as e:
        print(f"[Civitai Utils] Error getting local hashes: {e}")
        return web.json_response({"status": "error", "message": str(e)}, status=500)


@prompt_server.routes.get("/civitai_utils/get_scan_status")
async def get_scan_status(request):
    """用于查询后台扫描状态"""
    is_complete = utils.db_manager.get_setting("initial_scan_complete", False)
    return web.json_response({"status": "ok", "is_scanning": not is_complete})


# --- WebSocket 部分 ---
ACTIVE_CONNECTIONS = set()

@prompt_server.routes.get("/civitai_toolkit/ws")
async def websocket_handler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    ACTIVE_CONNECTIONS.add(ws)
    print("[Civitai Toolkit] New WebSocket connection established.")
    try:
        async for msg in ws:
            pass
    except Exception as e:
        print(f"[Civitai Toolkit] WebSocket connection error: {e}")
    finally:
        ACTIVE_CONNECTIONS.remove(ws)
        print("[Civitai Toolkit] WebSocket connection closed.")
    return ws

async def send_ws_message_async(msg_type, data):
    """ 异步发送函数 """
    message = json.dumps({"type": msg_type, "data": data})
    tasks = [conn.send_str(message) for conn in ACTIVE_CONNECTIONS]
    if tasks:
        await asyncio.gather(*tasks, return_exceptions=True)

def send_ws_message(msg_type, data):
    """
    [核心修改] 这是供后台线程调用的线程安全函数
    """
    if not main_loop or not main_loop.is_running():
        print("[Civitai Toolkit] Main event loop not available for WebSocket message.")
        return
    # 使用 run_coroutine_threadsafe 从同步线程向异步循环提交任务
    asyncio.run_coroutine_threadsafe(
        send_ws_message_async(msg_type, data),
        main_loop
    )