"""
trMemo Retrieval Tool - đọc ghi chú từ MongoDB cho Agent.

Dùng để trả lời các câu hỏi kiểu:
- "Hôm nay tôi đã note gì?"
- "Tuần trước tôi note những gì?"
- "Tìm memo có nội dung 'dự án A'" (dùng vector search)
- "Lọc memo theo tag #work"
"""

from __future__ import annotations

import json
import logging
import os
from datetime import datetime, time, timedelta, timezone
from typing import Optional, Any

from langchain_core.runnables import RunnableConfig

from langchain_core.tools import tool
from langchain_openai import OpenAIEmbeddings
from pydantic import SecretStr

from common.mongo_client import mongodb_client, serialize_doc


logger = logging.getLogger(__name__)

# ---- Cached singletons to avoid re-init per request ----
_embedder: OpenAIEmbeddings | None = None
_cached_embeddings: list[dict] | None = None
_cache_timestamp: float = 0
_CACHE_TTL = 300  # 5 minutes


def _get_embedder() -> OpenAIEmbeddings | None:
    global _embedder
    if _embedder is not None:
        return _embedder
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        return None
    _embedder = OpenAIEmbeddings(
        model="text-embedding-3-small",
        api_key=SecretStr(api_key),
    )
    return _embedder


async def _get_cached_embeddings() -> list[dict]:
    """Load embeddings from MongoDB with in-memory cache."""
    import time as _time
    global _cached_embeddings, _cache_timestamp
    now = _time.time()
    if _cached_embeddings is not None and (now - _cache_timestamp) < _CACHE_TTL:
        return _cached_embeddings
    cursor = mongodb_client.memo_embeddings.find({})
    _cached_embeddings = await cursor.to_list(length=5000)
    _cache_timestamp = now
    logger.info("Refreshed embeddings cache: %d docs", len(_cached_embeddings))
    return _cached_embeddings


@tool
async def memo_retrieval_tool(
    start_date: str,
    end_date: Optional[str] = None,
    content_search: Optional[str] = None,
    tag: Optional[str] = None,
    *,
    config: RunnableConfig,
) -> str:
    """
    Truy vấn các memo từ database MongoDB (cuccu_memos).

    Args:
        start_date: Ngày bắt đầu (YYYY-MM-DD). Bắt buộc.
        end_date: Ngày kết thúc (YYYY-MM-DD). Optional. Nếu không có thì mặc định tìm trong start_date.
        content_search: Từ khóa tìm kiếm trong nội dung (dùng Regex). Optional.
        tag: Lọc theo tag chính xác (ví dụ: "work", "idea"). Optional.

    Returns:
        JSON string chứa danh sách memo tìm được.
    """
    try:
        # Extract user_id from RunnableConfig to filter by creator
        user_id = (config.get("configurable") or {}).get("user_id")
        logger.info(
            "memo_retrieval_tool started: user_id=%s, start=%s, end=%s, content=%s, tag=%s",
            user_id, start_date, end_date, content_search, tag
        )

        # ──────────────────────────────────────────────
        # Step 1: Parse dates
        # ──────────────────────────────────────────────
        try:
            from common.timezone_config import VIETNAM_TZ
            dt_start = datetime.strptime(start_date, "%Y-%m-%d").replace(tzinfo=VIETNAM_TZ)
            dt_end = (
                datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=VIETNAM_TZ)
                if end_date else dt_start
            )
            begin_of_start_day = datetime.combine(dt_start, time.min).replace(tzinfo=VIETNAM_TZ).astimezone(timezone.utc)
            end_of_end_day = datetime.combine(dt_end, time.max).replace(tzinfo=VIETNAM_TZ).astimezone(timezone.utc)
        except ValueError as e:
            return json.dumps({
                "status": "error",
                "message": f"Invalid date format. Use YYYY-MM-DD. Error: {e}"
            }, ensure_ascii=False)

        # ──────────────────────────────────────────────
        # Step 2: Query MongoDB FIRST (date + tag filter)
        # → Narrow down memos before any semantic search
        # ──────────────────────────────────────────────
        base_query: dict[str, Any] = {
            "created_at": {
                "$gte": begin_of_start_day,
                "$lte": end_of_end_day,
            }
        }
        # Filter by creator_id — only return current user's memos
        if user_id:
            base_query["creator_id"] = user_id
        if tag:
            base_query["payload.tags"] = tag

        # If NO content_search → just run date/tag query directly
        if not content_search:
            logger.info("MongoDB Query (date-only): %s", base_query)
            cursor = mongodb_client.memos.find(base_query).sort("created_at", -1).limit(20)
            docs = await cursor.to_list(length=20)

        # If content_search with short keywords (≤2 words) → regex on MongoDB
        elif len(content_search.strip().split()) <= 2 or not _get_embedder():
            if not _get_embedder():
                logger.warning("⚠️ OPENAI_API_KEY not set, using regex search")
            logger.info(f"🔍 Using regex search for: '{content_search}'")
            base_query["content"] = {"$regex": content_search, "$options": "i"}
            logger.info("MongoDB Query (date+regex): %s", base_query)
            cursor = mongodb_client.memos.find(base_query).sort("created_at", -1).limit(20)
            docs = await cursor.to_list(length=20)

        # If content_search with longer queries → date query FIRST, then vector re-rank
        else:
            # Step 2a: Get all memos in date range (+ tag) from MongoDB
            logger.info("MongoDB Pre-filter Query: %s", base_query)
            cursor = mongodb_client.memos.find(base_query).sort("created_at", -1).limit(100)
            candidate_docs = await cursor.to_list(length=100)
            logger.info(f"📋 Pre-filter found {len(candidate_docs)} memos in date range")

            if not candidate_docs:
                docs = []
            else:
                # Step 2b: Get embeddings ONLY for these candidate memos
                candidate_ids = [str(doc["_id"]) for doc in candidate_docs]
                all_embeddings = await _get_cached_embeddings()
                # Filter embeddings to only candidates
                candidate_emb_map = {}
                for emb_doc in all_embeddings:
                    mid = str(emb_doc.get("memo_id", ""))
                    if mid in candidate_ids and emb_doc.get("embedding"):
                        candidate_emb_map[mid] = emb_doc["embedding"]

                logger.info(f"🧠 Found {len(candidate_emb_map)} embeddings for candidates")

                if not candidate_emb_map:
                    # No embeddings for these memos → fallback to regex
                    logger.info("No embeddings found, falling back to regex")
                    base_query["content"] = {"$regex": content_search, "$options": "i"}
                    cursor = mongodb_client.memos.find(base_query).sort("created_at", -1).limit(20)
                    docs = await cursor.to_list(length=20)
                else:
                    # Step 2c: Semantic re-rank using vector similarity
                    try:
                        embedder = _get_embedder()
                        query_embedding = await embedder.aembed_query(content_search)

                        import numpy as np
                        q = np.array(query_embedding, dtype=float)
                        scored: list[tuple[str, float]] = []

                        for mid, emb in candidate_emb_map.items():
                            v = np.array(emb, dtype=float)
                            if v.shape != q.shape:
                                continue
                            denom = np.linalg.norm(q) * np.linalg.norm(v)
                            sim = float(np.dot(q, v) / denom) if denom != 0 else 0.0
                            if sim > 0.4:
                                scored.append((mid, sim))

                        scored.sort(key=lambda x: x[1], reverse=True)
                        top_ids = {mid for mid, _ in scored[:20]}
                        logger.info(f"✅ Semantic re-rank: {len(top_ids)} memos above threshold")

                        # Re-order candidate_docs by semantic score
                        docs = [d for d in candidate_docs if str(d["_id"]) in top_ids]
                    except Exception as e:
                        logger.warning(f"⚠️ Vector search failed: {e}, falling back to regex")
                        base_query["content"] = {"$regex": content_search, "$options": "i"}
                        cursor = mongodb_client.memos.find(base_query).sort("created_at", -1).limit(20)
                        docs = await cursor.to_list(length=20)

        # ──────────────────────────────────────────────
        # Step 3: Format Result
        # ──────────────────────────────────────────────
        memos = [serialize_doc(doc) for doc in docs]

        return json.dumps(
            {
                "status": "success",
                "query": {
                    "start_date": start_date,
                    "end_date": end_date or start_date,
                    "content_search": content_search,
                    "tag": tag,
                },
                "count": len(memos),
                "memos": memos,
            },
            ensure_ascii=False,
            default=str,
        )

    except Exception as exc:
        logger.exception("memo_retrieval_tool error: %s", exc)
        return json.dumps(
            {
                "status": "error",
                "message": str(exc),
            },
            ensure_ascii=False,
        )
