# SPDX-License-Identifier: Apache-2.0
# Copyright 2025-2026 Raykosan (RaykoStudio)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import folder_paths
import json
import os
import struct
import hashlib
import re
import asyncio
import urllib.request
import urllib.error
from nodes import LoraLoader
from server import PromptServer
from datetime import datetime
import aiohttp

CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
PRESETS_DIR = os.path.join(CURRENT_DIR, "preset_models")
LORA_PRESETS_DIR = os.path.join(PRESETS_DIR, "loras")
RAYKO_LORA_DATA_DIR = os.path.join(CURRENT_DIR, "rayko_lora_data")

os.makedirs(PRESETS_DIR, exist_ok=True)
os.makedirs(LORA_PRESETS_DIR, exist_ok=True)
os.makedirs(RAYKO_LORA_DATA_DIR, exist_ok=True)


def clean_html(text):
    if not text:
        return ""
    clean = re.sub(r'<[^>]+>', ' ', str(text))
    clean = re.sub(r'\s+', ' ', clean).strip()
    return clean


def sanitize_preset_name(name):
    if not isinstance(name, str):
        return ""
    return "".join(c for c in name.strip() if c.isalnum() or c in " _-").strip()


def _safe_preset_path(name):
    safe_name = sanitize_preset_name(name)
    if not safe_name:
        return None
    base = os.path.realpath(LORA_PRESETS_DIR)
    candidate = os.path.realpath(os.path.join(base, f"{safe_name}.json"))
    try:
        if os.path.commonpath([base, candidate]) != base:
            return None
    except ValueError:
        return None
    return candidate


def _http_get_json(url, headers=None, timeout=30):
    req = urllib.request.Request(url, headers=headers or {})
    opener = urllib.request.build_opener()
    try:
        with opener.open(req, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8", errors="replace")
            try:
                return resp.status, json.loads(raw)
            except json.JSONDecodeError:
                return resp.status, None
    except urllib.error.HTTPError as e:
        return e.code, None


def extract_metadata_from_safetensors(lora_full_path):
    try:
        with open(lora_full_path, "rb") as f:
            header_size = struct.unpack("<Q", f.read(8))[0]
            header_bytes = f.read(header_size)
            header = json.loads(header_bytes.decode("utf-8"))
            if "__metadata__" in header:
                meta = header["__metadata__"]
                return {
                    "full_name": meta.get("modelspec.title", ""),
                    "trained_words": [t.strip() for t in meta.get("modelspec.tags", "").split(",") if t.strip()] if "modelspec.tags" in meta else [],
                    "description": clean_html(meta.get("modelspec.description", "") or meta.get("ss_training_comment", ""))
                }
    except Exception:
        pass
    return None


def compute_sha256(file_path):
    sha256_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            sha256_hash.update(chunk)
    return sha256_hash.hexdigest()


def save_to_rayko_db(file_hash, data, source, overwrite=False):
    db_path = os.path.join(RAYKO_LORA_DATA_DIR, f"{file_hash}.json")
    if not overwrite and os.path.exists(db_path):
        return False
    db_data = {
        "hash": file_hash,
        "name": data.get("full_name", ""),
        "trained_words": data.get("trained_words", []),
        "description": data.get("description", ""),
        "source": source,
        "imported_at": datetime.now().isoformat()
    }
    with open(db_path, 'w', encoding='utf-8') as f:
        json.dump(db_data, f, indent=2, ensure_ascii=False)
    return True


class RaykoLoRALoader:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL", {"tooltip": "Input model to apply LoRA"}),
                "use_clip": ("BOOLEAN", {"default": True, "label": "Apply LoRA to CLIP"}),
                "lora_data": ("STRING", {"default": "[]", "multiline": False, "forceInput": False}),
            },
            "optional": {
                "clip": ("CLIP", {"tooltip": "Input CLIP to apply LoRA (optional)"}),
            }
        }

    RETURN_TYPES = ("MODEL", "CLIP")
    RETURN_NAMES = ("MODEL", "CLIP")
    FUNCTION = "load_lora"
    CATEGORY = "🦊 RaykoStudio"
    DESCRIPTION = "Applies multiple LoRA models to the input model and optionally to CLIP."

    @classmethod
    def IS_CHANGED(cls, lora_data="[]", **kwargs):
        if lora_data is None:
            lora_data = "[]"
        return hashlib.md5(lora_data.encode()).hexdigest()

    def load_lora(self, model, use_clip, clip=None, lora_data="[]"):
        if not lora_data:
            lora_data = "[]"
        try:
            loras = json.loads(lora_data) if lora_data else []
            if not isinstance(loras, list):
                loras = []
        except Exception:
            loras = []

        lora_loader = LoraLoader()
        for lora in loras:
            if not isinstance(lora, dict):
                continue
            name = lora.get("name", "")
            strength_model = float(lora.get("strength_model", 1.0))
            strength_clip = float(lora.get("strength_clip", strength_model))
            enabled = lora.get("enabled", True)

            if name and name != "None" and enabled:
                lora_relative_path = name.replace("\\", "/")
                lora_full_path = folder_paths.get_full_path("loras", lora_relative_path)
                if lora_full_path:
                    try:
                        if use_clip and clip is not None:
                            model, clip = lora_loader.load_lora(model, clip, lora_relative_path, strength_model, strength_clip)
                        else:
                            model, _ = lora_loader.load_lora(model, None, lora_relative_path, strength_model, strength_clip)
                    except Exception as e:
                        print(f"[Rayko] LoRA Error: {e}")

        if use_clip and clip is not None:
            return (model, clip)
        else:
            return (model, None)


NODE_CLASS_MAPPINGS = {"RaykoLoRALoader": RaykoLoRALoader}
NODE_DISPLAY_NAME_MAPPINGS = {"RaykoLoRALoader": "🦊 RS LoRA Loader"}


@PromptServer.instance.routes.get("/rayko_lora_loader/get_loras")
async def get_loras(request):
    return aiohttp.web.json_response(sorted(folder_paths.get_filename_list("loras"), key=lambda x: x.lower()))


@PromptServer.instance.routes.post("/rayko_lora_loader/save_preset")
async def rayko_lora_loader_save_preset(request):
    try:
        data = await request.json()
        filepath = _safe_preset_path(data.get("name", ""))
        if not filepath:
            return aiohttp.web.Response(status=400, text="Invalid name")
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump({"lora_rows": data.get("lora_rows", [])}, f, indent=2)
        return aiohttp.web.Response(status=200, text="OK")
    except Exception as e:
        return aiohttp.web.Response(status=500, text=str(e))


@PromptServer.instance.routes.post("/rayko_lora_loader/list_presets")
async def rayko_lora_loader_list_presets(request):
    try:
        presets = [f[:-5] for f in os.listdir(LORA_PRESETS_DIR) if f.endswith('.json')] if os.path.exists(LORA_PRESETS_DIR) else []
        return aiohttp.web.json_response(sorted(presets, key=lambda x: x.lower()))
    except Exception as e:
        return aiohttp.web.Response(status=500, text=str(e))


@PromptServer.instance.routes.post("/rayko_lora_loader/load_preset")
async def rayko_lora_loader_load_preset(request):
    try:
        data = await request.json()
        filepath = _safe_preset_path(data.get("name", ""))
        if not filepath:
            return aiohttp.web.Response(status=400, text="Invalid name")
        if os.path.exists(filepath):
            with open(filepath, 'r', encoding='utf-8') as f:
                return aiohttp.web.json_response(json.load(f))
        return aiohttp.web.Response(status=404, text="Preset not found")
    except Exception as e:
        return aiohttp.web.Response(status=500, text=str(e))


@PromptServer.instance.routes.post("/rayko_lora_loader/delete_preset")
async def rayko_lora_loader_delete_preset(request):
    try:
        data = await request.json()
        filepath = _safe_preset_path(data.get("name", ""))
        if not filepath:
            return aiohttp.web.Response(status=400, text="Invalid name")
        if os.path.exists(filepath):
            os.remove(filepath)
            return aiohttp.web.Response(status=200, text="OK")
        return aiohttp.web.Response(status=404, text="Preset not found")
    except Exception as e:
        return aiohttp.web.Response(status=500, text=str(e))


@PromptServer.instance.routes.post("/rayko_lora_loader/get_lora_info")
async def rayko_lora_loader_get_lora_info(request):
    try:
        data = await request.json()
        lora_name = data.get("name", "")
        if not lora_name:
            return aiohttp.web.Response(status=400, text="No name provided")

        lora_relative_path = lora_name.replace("\\", "/")
        lora_full_path = folder_paths.get_full_path("loras", lora_relative_path)

        if not lora_full_path or not os.path.exists(lora_full_path):
            return aiohttp.web.json_response({"error": "File not found"})

        file_hash = compute_sha256(lora_full_path)

        db_path = os.path.join(RAYKO_LORA_DATA_DIR, f"{file_hash}.json")
        if os.path.exists(db_path):
            try:
                with open(db_path, 'r', encoding='utf-8') as f:
                    db_data = json.load(f)
                return aiohttp.web.json_response({
                    "full_name": db_data.get("name", lora_name),
                    "trained_words": db_data.get("trained_words", []),
                    "description": clean_html(db_data.get("description", "")),
                    "source": db_data.get("source", "rayko_db")
                })
            except Exception:
                pass

        safetensors_meta = extract_metadata_from_safetensors(lora_full_path)
        if safetensors_meta and (safetensors_meta["trained_words"] or safetensors_meta["description"]):
            save_to_rayko_db(file_hash, safetensors_meta, "safetensors")
            return aiohttp.web.json_response({
                "full_name": safetensors_meta["full_name"] or lora_name,
                "trained_words": safetensors_meta["trained_words"],
                "description": safetensors_meta["description"],
                "source": "safetensors"
            })

        return aiohttp.web.json_response({
            "full_name": lora_name,
            "trained_words": [],
            "description": "",
            "source": "none"
        })

    except Exception as e:
        return aiohttp.web.Response(status=500, text=str(e))


@PromptServer.instance.routes.post("/rayko_lora_loader/fetch_civitai_info")
async def rayko_lora_loader_fetch_civitai_info(request):
    try:
        data = await request.json()
        lora_name = data.get("name", "")
        if not lora_name:
            return aiohttp.web.Response(status=400, text="No name provided")

        lora_relative_path = lora_name.replace("\\", "/")
        lora_full_path = folder_paths.get_full_path("loras", lora_relative_path)

        if not lora_full_path or not os.path.exists(lora_full_path):
            return aiohttp.web.json_response({"error": "File not found"})

        file_hash = compute_sha256(lora_full_path)
        civitai_url = f"https://civitai.com/api/v1/model-versions/by-hash/{file_hash}"
        headers = {"User-Agent": "ComfyUI-RaykoStudio/1.0"}

        status, civitai_data = await asyncio.to_thread(_http_get_json, civitai_url, headers)

        if status == 404:
            return aiohttp.web.json_response({"error": "not_found", "message": "Not found on Civitai."})
        elif status == 429:
            return aiohttp.web.json_response({"error": "rate_limit", "message": "Rate limit reached."})
        elif status != 200 or not civitai_data:
            return aiohttp.web.json_response({"error": "api_error", "message": f"API status {status}"})

        model_id = civitai_data.get("modelId")
        model_tags = []

        if model_id:
            try:
                model_url = f"https://civitai.com/api/v1/models/{model_id}"
                m_status, model_data = await asyncio.to_thread(_http_get_json, model_url, headers)
                if m_status == 200 and isinstance(model_data, dict):
                    model_tags = model_data.get("tags", [])
            except Exception as e:
                print(f"[Rayko] Error fetching full model details: {e}")

        raw_trained = civitai_data.get("trainedWords", [])

        trained_words = []
        if isinstance(raw_trained, list):
            trained_words.extend([str(t).strip() for t in raw_trained if str(t).strip()])
        if isinstance(model_tags, list):
            trained_words.extend([str(t).strip() for t in model_tags if str(t).strip()])

        trained_words = list(dict.fromkeys(trained_words))

        raw_description = civitai_data.get("description", "")
        clean_description = clean_html(raw_description)

        model_name = civitai_data.get("model", {}).get("name", "")
        version_name = civitai_data.get("name", "")

        if model_name:
            if version_name and version_name.lower() not in ["v1", "v1.0", "default", "latest", ""]:
                full_name = f"{model_name} ({version_name})"
            else:
                full_name = model_name
        else:
            full_name = lora_name

        cache_data = {
            "full_name": full_name,
            "trained_words": trained_words,
            "description": clean_description
        }

        save_to_rayko_db(file_hash, cache_data, "civitai", overwrite=True)

        return aiohttp.web.json_response({
            "full_name": full_name,
            "trained_words": trained_words,
            "description": clean_description,
            "source": "civitai"
        })

    except urllib.error.URLError as e:
        return aiohttp.web.json_response({"error": "network", "message": f"Network error: {str(e)}"})
    except Exception as e:
        return aiohttp.web.Response(status=500, text=str(e))


@PromptServer.instance.routes.post("/rayko_lora_loader/add_tags")
async def rayko_add_tags(request):
    try:
        data = await request.json()
        lora_name = data.get("name", "")
        new_tags = data.get("tags", [])

        if not lora_name:
            return aiohttp.web.Response(status=400, text="No name provided")
        if not isinstance(new_tags, list) or not new_tags:
            return aiohttp.web.json_response({"error": "No tags provided"})

        lora_relative_path = lora_name.replace("\\", "/")
        lora_full_path = folder_paths.get_full_path("loras", lora_relative_path)

        if not lora_full_path or not os.path.exists(lora_full_path):
            return aiohttp.web.json_response({"error": "File not found"})

        file_hash = compute_sha256(lora_full_path)
        db_path = os.path.join(RAYKO_LORA_DATA_DIR, f"{file_hash}.json")

        existing_data = {}
        if os.path.exists(db_path):
            try:
                with open(db_path, 'r', encoding='utf-8') as f:
                    existing_data = json.load(f)
            except Exception:
                pass

        existing_tags = existing_data.get("trained_words", [])
        if not isinstance(existing_tags, list):
            existing_tags = []

        combined_tags = existing_tags.copy()
        for tag in new_tags:
            tag = str(tag).strip()
            if tag and tag not in combined_tags:
                combined_tags.append(tag)

        save_data = {
            "full_name": existing_data.get("name", ""),
            "trained_words": combined_tags,
            "description": existing_data.get("description", "")
        }

        save_to_rayko_db(file_hash, save_data, "manual", overwrite=True)

        return aiohttp.web.json_response({
            "trained_words": combined_tags,
            "added": len(combined_tags) - len(existing_tags)
        })

    except Exception as e:
        return aiohttp.web.Response(status=500, text=str(e))


@PromptServer.instance.routes.post("/rayko_lora_loader/remove_tag")
async def rayko_remove_tag(request):
    try:
        data = await request.json()
        lora_name = data.get("name", "")
        tag_to_remove = data.get("tag", "")

        if not lora_name or not tag_to_remove:
            return aiohttp.web.Response(status=400, text="Missing parameters")

        lora_relative_path = lora_name.replace("\\", "/")
        lora_full_path = folder_paths.get_full_path("loras", lora_relative_path)

        if not lora_full_path or not os.path.exists(lora_full_path):
            return aiohttp.web.json_response({"error": "File not found"})

        file_hash = compute_sha256(lora_full_path)
        db_path = os.path.join(RAYKO_LORA_DATA_DIR, f"{file_hash}.json")

        if not os.path.exists(db_path):
            return aiohttp.web.json_response({"error": "No data found"})

        with open(db_path, 'r', encoding='utf-8') as f:
            existing_data = json.load(f)

        existing_tags = existing_data.get("trained_words", [])
        if tag_to_remove in existing_tags:
            existing_tags.remove(tag_to_remove)

        save_data = {
            "full_name": existing_data.get("name", ""),
            "trained_words": existing_tags,
            "description": existing_data.get("description", "")
        }

        save_to_rayko_db(file_hash, save_data, existing_data.get("source", "manual"), overwrite=True)

        return aiohttp.web.json_response({
            "trained_words": existing_tags,
            "removed": tag_to_remove
        })

    except Exception as e:
        return aiohttp.web.Response(status=500, text=str(e))