# -*- coding: utf-8 -*-
"""
로라 파일 하나를 "LLM에게 설명할 수 있는 프로필"로 변환한다.

기존 `get_triggers()`(dolphin_vision_director.py)는 트리거 '단어'만 뽑아서
캐릭터 이름 옆에 콤마로 나열하는 용도였다. 그 방식은 문법적 맥락이 없어서
`bounce`/`swing` 같은 범용 동작 동사가 전신 모션으로 새는 문제가 있었고,
대응으로 그런 단어를 아예 제외하는 필터(_filter_generic_motion_verbs)를
넣어둔 상태다 - 로라가 의도한 효과까지 같이 사라지는 우회책이다.

이 모듈은 다른 방향으로 접근한다: 트리거를 '나중에 덧붙일 토큰'이 아니라
'프롬프트를 쓰기 전에 LLM에게 줄 재료'로 본다. LoRA Manager가 파일 옆에
만들어두는 `<이름>.metadata.json` 사이드카에 이미 다음이 다 들어있다.

  - model_name / civitai.model.name : 이 로라의 이름
  - modelDescription               : 자연어 설명(HTML)
  - tags / civitai.model.tags      : 분류 태그
  - civitai.trainedWords           : 학습 트리거. 단어인 경우도 있지만
                                     상당수는 학습 캡션 문장 그대로다
                                     (예: "She pulls down her pants revealing a penis.")
  - base_model                     : Wan 2.1/2.2, i2v/t2v 구분

로라는 학습 캡션 분포와 닮은 프롬프트에서 가장 잘 발화하므로, 이 재료를
LLM에 먼저 주면 트리거를 문장 안에 문법적으로 녹인 프롬프트를 쓸 수 있다.

trainedWords 3분류 (`classify_trained_words` 참고):
  hard  : 공백 없고 숫자나 대문자가 섞인 학습 식별자. LLM이 "고쳐주면"
          로라가 안 걸리므로 원문 그대로 나와야 한다 (예: `4n4lfi5tin6`,
          `PENISLORA`). 출력 검증 대상.
  soft  : 공백 없는 평범한 소문자 영단어 (예: `bounce`, `penis`).
          강제하지 않는다 - 강제 삽입이 바로 예전 버그의 원인이었다.
          문맥에 맞을 때 몸 부위에 붙여 쓰라고 '권고'만 한다.
  caption: 공백이 있는 구/문장. 학습 캡션 예시로만 준다(문체 참고용).
"""
import json
import os
import re

_HTML_TAG_RE = re.compile(r"<[^>]+>")
_WS_RE = re.compile(r"\s+")

# 세션 내 프로필 캐시. 키는 로라 파일의 (경로, size, mtime).
_PROFILE_CACHE = {}

# LLM 컨텍스트를 잡아먹지 않도록 설명은 잘라서 준다. 로라 설명문은 보통
# 앞부분에 "무엇을 하는 로라인지"가 있고 뒤는 학습 설정/후원 링크다.
MAX_DESC_CHARS = 400
MAX_CAPTION_EXAMPLES = 2


def _strip_html(s):
    if not s:
        return ""
    s = re.sub(r"<br\s*/?>", " ", s, flags=re.IGNORECASE)
    s = re.sub(r"</p>", " ", s, flags=re.IGNORECASE)
    s = _HTML_TAG_RE.sub("", s)
    s = (s.replace("&nbsp;", " ").replace("&amp;", "&")
          .replace("&lt;", "<").replace("&gt;", ">").replace("&quot;", '"'))
    return _WS_RE.sub(" ", s).strip()


def _truncate(s, limit):
    if not s or len(s) <= limit:
        return s
    cut = s[:limit]
    # 문장 중간에서 자르지 않도록 마지막 문장부호까지만.
    for sep in (". ", "! ", "? "):
        idx = cut.rfind(sep)
        if idx > limit * 0.5:
            return cut[:idx + 1]
    return cut.rstrip() + "..."


def is_hard_token(word):
    """LLM이 철자를 바꾸면 안 되는 학습 식별자인가.

    공백이 없고(=단일 토큰) 숫자나 대문자가 섞여있으면 학습 식별자로 본다.
    `4n4lfi5tin6`, `PENISLORA`, `az420` -> True
    `bounce`, `penis`, `swing`          -> False (평범한 소문자 영단어)
    """
    w = word.strip()
    if not w or _WS_RE.search(w):
        return False
    return any(c.isdigit() for c in w) or any(c.isupper() for c in w)


def classify_trained_words(words):
    """trainedWords 리스트를 (hard, soft, captions)로 나눈다."""
    hard, soft, captions = [], [], []
    for raw in words or []:
        w = str(raw).strip().strip(",")
        if not w:
            continue
        if _WS_RE.search(w):
            captions.append(w)
        elif is_hard_token(w):
            hard.append(w)
        else:
            soft.append(w)
    # 순서 유지 중복 제거
    return (list(dict.fromkeys(hard)), list(dict.fromkeys(soft)),
            list(dict.fromkeys(captions)))


def _sidecar_path(lora_path):
    """LoRA Manager가 만드는 `<확장자 뗀 이름>.metadata.json`."""
    return os.path.splitext(lora_path)[0] + ".metadata.json"


def _read_sidecar(lora_path):
    p = _sidecar_path(lora_path)
    if not os.path.exists(p):
        return None
    try:
        with open(p, "r", encoding="utf-8") as f:
            return json.load(f)
    except (OSError, ValueError, json.JSONDecodeError) as e:
        print(f"⚠️ [LoraProfile] 사이드카 파싱 실패({os.path.basename(p)}): "
              f"{type(e).__name__}: {e}")
        return None


def _resolve_lora_path(lora_name):
    try:
        import folder_paths
        return folder_paths.get_full_path("loras", lora_name)
    except Exception:
        return lora_name if os.path.exists(lora_name) else None


def build_profile_from_sidecar(data, lora_name, lora_path):
    """사이드카 dict -> 프로필 dict. 사이드카가 없으면 호출하지 않는다."""
    civitai = data.get("civitai") or {}
    model = civitai.get("model") or {}

    name = (data.get("model_name") or model.get("name")
            or os.path.splitext(os.path.basename(lora_name))[0])
    desc = _strip_html(data.get("modelDescription") or model.get("description") or "")
    tags = data.get("tags") or model.get("tags") or []
    if not isinstance(tags, list):
        tags = []
    trained = civitai.get("trainedWords") or []
    if not isinstance(trained, list):
        trained = [trained]
    hard, soft, captions = classify_trained_words(trained)

    return {
        "lora_name": lora_name,
        "file": os.path.basename(lora_name),
        "display_name": str(name).strip(),
        "description": _truncate(desc, MAX_DESC_CHARS),
        "tags": [str(t).strip() for t in tags if str(t).strip()],
        "base_model": str(data.get("base_model") or civitai.get("baseModel") or "").strip(),
        "hard_tokens": hard,
        "soft_words": soft,
        "captions": captions[:MAX_CAPTION_EXAMPLES],
        "source": "sidecar",
    }


def load_profile(lora_name, fallback_triggers=None):
    """로라 이름(ComfyUI 상대경로) -> 프로필 dict. 실패 시 None.

    사이드카가 없으면 `fallback_triggers`(기존 get_triggers 경로의 결과)로
    최소 프로필이라도 만든다 - 개인 학습 로라처럼 civitai에 없는 경우.

    `fallback_triggers`는 리스트 또는 "호출하면 리스트를 주는 함수"를 받는다.
    기존 조회 경로는 sha256 해시(수 GB 파일) + 네트워크 호출까지 할 수 있어서,
    사이드카가 있는 99.5%의 경우엔 아예 호출되지 않도록 지연 평가한다.
    """
    if not lora_name or lora_name == "None":
        return None

    path = _resolve_lora_path(lora_name)
    cache_key = lora_name
    if path and os.path.exists(path):
        try:
            st = os.stat(path)
            cache_key = f"{lora_name}::{st.st_size}::{int(st.st_mtime)}"
        except OSError:
            pass
    if cache_key in _PROFILE_CACHE:
        return _PROFILE_CACHE[cache_key]

    profile = None
    if path:
        data = _read_sidecar(path)
        if data:
            profile = build_profile_from_sidecar(data, lora_name, path)

    if profile is None:
        # 사이드카 없음 - 기존 트리거 조회 결과로 최소 프로필.
        fb = fallback_triggers() if callable(fallback_triggers) else fallback_triggers
        hard, soft, captions = classify_trained_words(fb or [])
        if not (hard or soft or captions):
            return None
        profile = {
            "lora_name": lora_name,
            "file": os.path.basename(lora_name),
            "display_name": os.path.splitext(os.path.basename(lora_name))[0],
            "description": "",
            "tags": [],
            "base_model": "",
            "hard_tokens": hard,
            "soft_words": soft,
            "captions": captions[:MAX_CAPTION_EXAMPLES],
            "source": "triggers-only",
        }

    _PROFILE_CACHE[cache_key] = profile
    return profile


def format_profiles_for_llm(profiles, indent="  "):
    """프로필 목록 -> LLM 프롬프트에 넣을 텍스트 블록.

    한 로라당 4줄 이내로 압축한다. 로라가 4개까지 걸릴 수 있고 클립도 4개라
    장황하면 시스템 프롬프트의 다른 규칙이 묻힌다.
    """
    if not profiles:
        return ""
    lines = []
    for p in profiles:
        head = f"{indent}- {p['display_name']}"
        if p.get("tags"):
            head += f"  [tags: {', '.join(p['tags'][:5])}]"
        lines.append(head)
        if p.get("description"):
            lines.append(f"{indent}  what it does: {p['description']}")
        if p.get("hard_tokens"):
            lines.append(f"{indent}  MUST APPEAR VERBATIM: "
                         + ", ".join(f'"{t}"' for t in p["hard_tokens"]))
        if p.get("soft_words"):
            lines.append(f"{indent}  vocabulary it was trained on (use only where it "
                         f"genuinely fits, anchored to the body part): "
                         + ", ".join(p["soft_words"]))
        for cap in p.get("captions", []):
            lines.append(f'{indent}  training caption example: "{cap}"')
    return "\n".join(lines)


def collect_hard_tokens(profiles):
    out = []
    for p in profiles or []:
        out.extend(p.get("hard_tokens") or [])
    return list(dict.fromkeys(out))


def missing_hard_tokens(text, profiles):
    """생성된 프롬프트에 빠진 hard token 목록.

    LLM은 `4n4lfi5tin6` 같은 토큰을 "정상 철자로 고쳐주는" 경향이 있어서
    (그 순간 로라가 안 걸린다) 대소문자 무시하고 존재만 확인한다.
    """
    low = (text or "").lower()
    return [t for t in collect_hard_tokens(profiles) if t.lower() not in low]


if __name__ == "__main__":
    # 실제 로라 폴더에 대고 돌려보는 자체 점검.
    import glob
    import sys

    # ComfyUI 콘솔은 UTF-8이지만 맨손으로 돌릴 땐 cp949라 이모지 print가 죽는다.
    try:
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    except (AttributeError, OSError):
        pass

    root = sys.argv[1] if len(sys.argv) > 1 else "."
    files = sorted(glob.glob(os.path.join(root, "*.safetensors")))
    print(f"safetensors {len(files)}개 검사\n")
    stats = {"sidecar": 0, "none": 0, "hard": 0, "soft": 0, "caption": 0, "desc": 0}
    samples = []
    for f in files:
        data = _read_sidecar(f)
        if not data:
            stats["none"] += 1
            continue
        p = build_profile_from_sidecar(data, os.path.basename(f), f)
        stats["sidecar"] += 1
        if p["hard_tokens"]:
            stats["hard"] += 1
        if p["soft_words"]:
            stats["soft"] += 1
        if p["captions"]:
            stats["caption"] += 1
        if p["description"]:
            stats["desc"] += 1
        if p["hard_tokens"] and len(samples) < 5:
            samples.append(p)
    print(f"사이드카 있음: {stats['sidecar']} / 없음: {stats['none']}")
    print(f"  설명 있음:        {stats['desc']}")
    print(f"  hard token 있음:  {stats['hard']}")
    print(f"  soft word 있음:   {stats['soft']}")
    print(f"  caption 예시 있음: {stats['caption']}\n")
    for p in samples:
        print("=" * 60)
        print(format_profiles_for_llm([p]))