import os
import re

from huggingface_hub import InferenceClient

# the functions for resolving and generating titles
# it tries to query and hf model for title

# some shitty fallback logic, if models fail
# could improve candidate model defining code

def title_from_query(query: str) -> str:
    stop_words = {
        "a", "an", "and", "are", "as", "at", "be", "by", "can", "do", "for", "from", "how",
        "i", "in", "is", "it", "me", "my", "of", "on", "or", "please", "show", "tell", "that",
        "the", "this", "to", "we", "what", "when", "where", "which", "why", "with", "you", "your",
    }

    words = re.findall(r"[A-Za-z0-9][A-Za-z0-9\-_/+]*", query)
    if not words:
        return "New Chat"

    filtered: list[str] = []
    for word in words:
        cleaned = word.strip("-_/+")
        if not cleaned:
            continue
        if cleaned.lower() in stop_words:
            continue
        filtered.append(cleaned)
        if len(filtered) >= 6:
            break

    chosen = filtered if filtered else words[:6]
    normalized = [w.capitalize() if w.islower() else w for w in chosen]
    title = " ".join(normalized).strip()
    return title[:80] if title else "New Chat"


def clean_title_text(raw: str) -> str:
    text = (raw or "").strip()
    text = text.replace("\n", " ").replace("\r", " ")
    text = re.sub(r"^[\"'`\s]+|[\"'`\s]+$", "", text)
    text = re.sub(r"\s+", " ", text).strip()
    words = text.split()
    if len(words) > 8:
        text = " ".join(words[:8])
    return text[:80]


def title_from_hf(query: str, client: InferenceClient, model_id: str) -> str | None:
    system_prompt = (
        "You generate short chat titles. Return only a title, no punctuation at the end, no quotes."
    )
    user_prompt = (
        "Create a concise 3-7 word title for this user request:\n"
        f"{query}"
    )

    response = client.chat_completion(
        model=model_id,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        max_tokens=24,
        temperature=0.3,
    )
    if not response or not response.choices:
        return None

    raw_title = response.choices[0].message.content or ""
    title = clean_title_text(raw_title)
    if not title or title.lower() == "new chat":
        return None
    return title


def parse_title_model_candidates() -> list[str]:
    
    raw = os.getenv(
        "TITLE_MODEL_IDS",
        "Qwen/Qwen2.5-1.5B-Instruct,CohereLabs/tiny-aya-global,meta-llama/Meta-Llama-3-8B-Instruct",
    )
    models = [m.strip() for m in raw.split(",") if m.strip()]
    return models or ["meta-llama/Meta-Llama-3-8B-Instruct"]