#!/usr/bin/env python3
"""
Гибридный поиск: Meilisearch + ChromaDB + LLM суммаризация.
"""

import json
import os
import re
import sys
import time
from pathlib import Path
from typing import Optional

import requests as _requests

# Сессия без прокси для внутренних сервисов (Meilisearch/ChromaDB на localhost)
_internal_session = _requests.Session()
_internal_session.trust_env = False

# Загружаем .env
def load_env():
    env_file = Path.home() / ".openclaw" / ".env"
    if env_file.exists():
        with open(env_file) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, _, val = line.partition("=")
                    os.environ.setdefault(key.strip(), val.strip())

load_env()

MEILI_URL = "http://127.0.0.1:7700"
CHROMA_PATH = str(Path(__file__).parent.parent / "data" / "chromadb")
COLLECTION_NAME = "snowbike_embeddings"
INDEX_NAME = "snowbike_messages"
MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
OPENROUTER_KEY = os.environ.get("OPENROUTER_API_KEY", "")
LLM_MODEL = "anthropic/claude-sonnet-4-5"

# Порог релевантности для ChromaDB (меньше = лучше совпадение)
CHROMA_DISTANCE_THRESHOLD = 1.2

# === СЕЗОННЫЙ ФИЛЬТР ===
# Маппинг ключевых слов → месяцы (включая соседние для контекста)
SEASON_KEYWORDS = {
    # Месяцы
    "январ": [12, 1, 2],
    "феврал": [1, 2, 3],
    "март": [2, 3, 4],
    "апрел": [3, 4, 5],
    "май": [4, 5, 6],
    "июн": [5, 6, 7],
    "июл": [6, 7, 8],
    "август": [7, 8, 9],
    "сентябр": [8, 9, 10],
    "октябр": [9, 10, 11],
    "ноябр": [10, 11, 12],
    "декабр": [11, 12, 1],
    # Сезоны
    "зим": [12, 1, 2],
    "весн": [3, 4, 5],
    "лет": [6, 7, 8],
    "осен": [9, 10, 11],
    # Сезон катания
    "сезон": [11, 12, 1, 2, 3, 4],  # типичный сноубайк-сезон
    "межсезон": [5, 6, 7, 8, 9, 10],
}


def detect_season_months(query: str) -> Optional[list]:
    """
    Определяет месяцы из запроса.
    Возвращает список номеров месяцев (1-12) или None.
    """
    query_lower = query.lower()
    for keyword, months in SEASON_KEYWORDS.items():
        if keyword in query_lower:
            return months
    return None


SYSTEM_PROMPT = """Ты — помощник по сноубайкам. Отвечаешь на вопросы на основе сообщений из Telegram-группы «Сноубайк Россия».

Правила:
1. Отвечай только на русском языке
2. Используй только информацию из предоставленных сообщений
3. Если информации недостаточно — честно скажи об этом, перечисли что было упомянуто
4. Всегда указывай конкретику: ДАТЫ (в формате «в апреле 2024», «в январе 2023» и т.д.), имена (если есть), цифры
5. Агрегируй мнения: если несколько человек говорят одно и то же — обобщи
6. Будь конкретен: перечисли проблемы/решения по пунктам
7. При ответе на сезонные вопросы — обязательно указывай из каких месяцев/лет взяты данные"""

USER_PROMPT_TEMPLATE = """Вопрос: {question}

Найденные сообщения из чата «Сноубайк Россия»:
{context}

Дай связный ответ на вопрос. Если нашлись конкретные факты — перечисли их. 
Обязательно упоминай даты сообщений (месяц и год) когда это важно для контекста."""


# Глобальные клиенты (инициализируются лениво)
_chroma_collection = None
_embed_model = None


def get_chroma_collection():
    global _chroma_collection
    if _chroma_collection is None:
        import chromadb
        client = chromadb.PersistentClient(path=CHROMA_PATH)
        _chroma_collection = client.get_collection(COLLECTION_NAME)
    return _chroma_collection


def get_embed_model():
    global _embed_model
    if _embed_model is None:
        from sentence_transformers import SentenceTransformer
        _embed_model = SentenceTransformer(MODEL_NAME)
    return _embed_model


def search_meilisearch(query: str, topic_ids: Optional[list] = None, limit: int = 20,
                       season_months: Optional[list] = None) -> list:
    """Полнотекстовый поиск в Meilisearch через HTTP."""
    try:
        fetch_limit = limit * 6 if season_months else limit
        params = {"q": query, "limit": fetch_limit}

        if topic_ids:
            filter_str = " OR ".join(f"topic_id = {t}" for t in topic_ids)
            params["filter"] = filter_str

        resp = _internal_session.post(
            f"{MEILI_URL}/indexes/{INDEX_NAME}/search",
            json=params,
            timeout=10,
        )
        resp.raise_for_status()
        hits = resp.json().get("hits", [])

        if season_months:
            # Разделяем на "в сезон" и "вне сезона", возвращаем сначала сезонные
            in_season = []
            out_season = []
            for hit in hits:
                date_str = hit.get("date", "")
                try:
                    month = int(date_str[5:7]) if len(date_str) >= 7 else 0
                except (ValueError, IndexError):
                    month = 0
                if month in season_months:
                    in_season.append(hit)
                else:
                    out_season.append(hit)
            hits = (in_season + out_season)[:limit]

        return hits[:limit]
    except Exception as e:
        print(f"Meilisearch ошибка: {e}", file=sys.stderr)
        return []


def search_chromadb(query: str, topic_ids: Optional[list] = None, limit: int = 20,
                    season_months: Optional[list] = None) -> list:
    """Семантический поиск в ChromaDB с фильтрацией по дистанции."""
    try:
        model = get_embed_model()
        collection = get_chroma_collection()

        embedding = model.encode(query).tolist()

        # Строим where-фильтр
        conditions = []

        if topic_ids:
            if len(topic_ids) == 1:
                conditions.append({"topic_id": topic_ids[0]})
            else:
                conditions.append({"$or": [{"topic_id": t} for t in topic_ids]})

        if season_months:
            # Используем поле month если доступно (после переиндексации)
            # Для старых записей без month — пробуем с запасом и пост-фильтрацией
            has_month_field = _check_month_field(collection)
            if has_month_field:
                if len(season_months) == 1:
                    conditions.append({"month": season_months[0]})
                else:
                    conditions.append({"$or": [{"month": m} for m in season_months]})

        where = None
        if len(conditions) == 1:
            where = conditions[0]
        elif len(conditions) > 1:
            where = {"$and": conditions}

        fetch_n = limit if not (season_months and not _check_month_field(collection)) else limit * 4

        kwargs = {
            "query_embeddings": [embedding],
            "n_results": min(fetch_n, collection.count()),
            "include": ["documents", "metadatas", "distances"],
        }
        if where:
            kwargs["where"] = where

        result = collection.query(**kwargs)

        hits = []
        ids = result.get("ids", [[]])[0]
        docs = result.get("documents", [[]])[0]
        metas = result.get("metadatas", [[]])[0]
        dists = result.get("distances", [[]])[0]

        for msg_id, doc, meta, dist in zip(ids, docs, metas, dists):
            if dist > CHROMA_DISTANCE_THRESHOLD:
                continue
            # Пост-фильтр по месяцу если нет поля month
            if season_months and not _check_month_field(collection):
                date_str = meta.get("date", "")
                try:
                    month = int(date_str[5:7]) if len(date_str) >= 7 else 0
                except (ValueError, IndexError):
                    month = 0
                if month not in season_months:
                    continue
            hits.append({
                "id": int(msg_id),
                "text": doc,
                "topic_id": meta.get("topic_id"),
                "topic_title": meta.get("topic_title", ""),
                "date": meta.get("date", ""),
                "from_id": meta.get("from_id", ""),
                "_chroma_distance": dist,
            })

        return hits[:limit]
    except Exception as e:
        print(f"ChromaDB ошибка: {e}", file=sys.stderr)
        return []


_month_field_cache = None

def _check_month_field(collection) -> bool:
    """Проверяем есть ли поле month в коллекции (кэшируем результат)."""
    global _month_field_cache
    if _month_field_cache is not None:
        return _month_field_cache
    try:
        sample = collection.get(limit=1, include=["metadatas"])
        metas = sample.get("metadatas", [])
        if metas and "month" in metas[0]:
            _month_field_cache = True
        else:
            _month_field_cache = False
    except Exception:
        _month_field_cache = False
    return _month_field_cache


def merge_results(meili_hits: list, chroma_hits: list, limit: int = 15) -> list:
    """Объединяем и дедуплицируем результаты."""
    seen_ids = set()
    merged = []

    # Meilisearch результаты — полнотекстовые, высокий приоритет
    for hit in meili_hits:
        msg_id = hit.get("id")
        if msg_id and msg_id not in seen_ids:
            hit["_source"] = "meilisearch"
            merged.append(hit)
            seen_ids.add(msg_id)

    # ChromaDB результаты — семантические, дополняем
    for hit in chroma_hits:
        msg_id = hit.get("id")
        if msg_id and msg_id not in seen_ids:
            hit["_source"] = "chromadb"
            merged.append(hit)
            seen_ids.add(msg_id)

    return merged[:limit]


def format_context(messages: list) -> str:
    """Форматируем контекст для LLM с явными датами."""
    lines = []
    for i, msg in enumerate(messages, 1):
        date_full = msg.get("date", "")
        # Форматируем дату как "15 апреля 2024"
        date_display = _format_date_russian(date_full)
        topic = msg.get("topic_title", "?")
        author = msg.get("from_id", "?")
        text = msg.get("text", "")
        msg_id = msg.get("id", "?")
        source = msg.get("_source", "")

        lines.append(f"[{i}] ID:{msg_id} | Дата: {date_display} | Топик: {topic} | автор:{author} | источник:{source}")
        lines.append(f"    {text}")
        lines.append("")

    return "\n".join(lines)


MONTH_NAMES_RU = {
    1: "января", 2: "февраля", 3: "марта", 4: "апреля",
    5: "мая", 6: "июня", 7: "июля", 8: "августа",
    9: "сентября", 10: "октября", 11: "ноября", 12: "декабря",
}


def _format_date_russian(date_str: str) -> str:
    """Преобразует ISO дату в читаемый русский формат."""
    try:
        # "2024-04-15T..."
        year = int(date_str[:4])
        month = int(date_str[5:7])
        day = int(date_str[8:10])
        return f"{day} {MONTH_NAMES_RU[month]} {year}"
    except Exception:
        return date_str[:10] if date_str else "?"


def ask_llm(question: str, context: str) -> str:
    """Задаём вопрос LLM через OpenRouter."""
    if not OPENROUTER_KEY:
        return "ОШИБКА: OPENROUTER_API_KEY не задан в ~/.openclaw/.env"

    user_msg = USER_PROMPT_TEMPLATE.format(question=question, context=context)

    response = _requests.post(
        "https://openrouter.ai/api/v1/chat/completions",
        headers={
            "Authorization": f"Bearer {OPENROUTER_KEY}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://snowbike-rag.local",
            "X-Title": "Snowbike RAG",
        },
        json={
            "model": LLM_MODEL,
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_msg},
            ],
            "max_tokens": 1500,
            "temperature": 0.3,
        },
        timeout=30,
    )

    response.raise_for_status()
    data = response.json()
    return data["choices"][0]["message"]["content"]


def search(query: str, topic_ids: Optional[list] = None, limit: int = 15) -> dict:
    """
    Главная функция гибридного поиска.

    Returns:
        dict с полями: query, answer, sources, count, time_ms
    """
    start_time = time.time()

    # Определяем сезонные месяцы
    season_months = detect_season_months(query)
    if season_months:
        print(f"Сезонный запрос, фильтр по месяцам: {season_months}", file=sys.stderr)

    # Поиск
    meili_hits = search_meilisearch(query, topic_ids, limit=20, season_months=season_months)
    chroma_hits = search_chromadb(query, topic_ids, limit=20, season_months=season_months)

    print(f"Meilisearch hits: {len(meili_hits)}, ChromaDB hits: {len(chroma_hits)}", file=sys.stderr)

    merged = merge_results(meili_hits, chroma_hits, limit=limit)

    if not merged:
        return {
            "query": query,
            "answer": "К сожалению, по вашему запросу ничего не найдено в базе сообщений.",
            "sources": [],
            "count": 0,
            "time_ms": int((time.time() - start_time) * 1000),
        }

    # Формируем контекст для LLM
    context = format_context(merged)

    # LLM ответ
    answer = ask_llm(query, context)

    # Источники
    sources = [
        {
            "id": msg.get("id"),
            "date": msg.get("date", "")[:10],
            "date_display": _format_date_russian(msg.get("date", "")),
            "topic": msg.get("topic_title", ""),
            "author": str(msg.get("from_id", "")),
            "text_preview": msg.get("text", "")[:120],
        }
        for msg in merged
    ]

    elapsed_ms = int((time.time() - start_time) * 1000)

    return {
        "query": query,
        "answer": answer,
        "sources": sources,
        "count": len(merged),
        "time_ms": elapsed_ms,
        "season_months": season_months,
    }


if __name__ == "__main__":
    query = sys.argv[1] if len(sys.argv) > 1 else "проблемы ремень трикстер"

    print(f"Запрос: {query}\n")
    result = search(query)

    print(f"Найдено источников: {result['count']}")
    print(f"Время: {result['time_ms']} мс")
    if result.get("season_months"):
        print(f"Сезонный фильтр (месяцы): {result['season_months']}")
    print("\n=== ОТВЕТ ===")
    print(result["answer"])
    print("\n=== ИСТОЧНИКИ ===")
    for s in result["sources"][:5]:
        print(f"  [{s['id']}] {s['date_display']} | {s['topic']} | {s['text_preview'][:100]}...")