"""
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.tools import tool
from langchain_openai import OpenAIEmbeddings
from pydantic import SecretStr
from bson import ObjectId

from common.mongo_client import mongodb_client, serialize_doc


logger = logging.getLogger(__name__)


@tool
async def memo_retrieval_tool(
    start_date: str,
    end_date: Optional[str] = None,
    content_search: Optional[str] = None,
    tag: Optional[str] = None,
) -> 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:
        logger.info(
            "memo_retrieval_tool started: start=%s, end=%s, content=%s, tag=%s",
            start_date, end_date, content_search, tag
        )

        # 1. Build Date Filter
        query: dict[str, Any] = {}
        
        try:
            # Parse start_date key
            dt_start = datetime.strptime(start_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
            
            if end_date:
                # Range query: start_date 00:00:00 -> end_date 23:59:59
                dt_end_str = end_date
                dt_end = datetime.strptime(end_date, "%Y-%m-%d").replace(tzinfo=timezone.utc)
            else:
                # Single day query: start_date 00:00:00 -> start_date 23:59:59
                dt_end_str = start_date
                dt_end = dt_start

            # Tạo mốc thời gian đầu ngày và cuối ngày
            # Lưu ý: MongoDB lưu UTC. Cần đảm bảo logic query phù hợp.
            # Ở đây giả sử start_date nhận vào là Local nhưng backend lưu UTC.
            # Tạm thời query range rộng (start 00:00 -> end 23:59:59)
            
            # Helper to get begin of day and end of day
            begin_of_start_day = datetime.combine(dt_start, time.min).replace(tzinfo=timezone.utc)
            end_of_end_day = datetime.combine(dt_end, time.max).replace(tzinfo=timezone.utc)

            query["created_at"] = {
                "$gte": begin_of_start_day,
                "$lte": end_of_end_day
            }

        except ValueError as e:
            return json.dumps({
                "status": "error",
                "message": f"Invalid date format. Use YYYY-MM-DD. Error: {e}"
            }, ensure_ascii=False)

        # 2. Content Search - Use Vector Search if content_search provided
        memo_ids_from_vector_search = []
        if content_search:
            # Use vector search for semantic understanding
            try:
                api_key = os.getenv("OPENAI_API_KEY")
                if api_key:
                    logger.info(f"🔍 Using vector search for: '{content_search}'")
                    embedder = OpenAIEmbeddings(
                        model="text-embedding-3-small",
                        api_key=SecretStr(api_key),
                    )
                    # Use async embedding to avoid blocking the event loop
                    query_embedding = await embedder.aembed_query(content_search)
                    
                    # Search embeddings collection
                    import numpy as np
                    cursor_emb = mongodb_client.memo_embeddings.find({})
                    emb_docs = await cursor_emb.to_list(length=1000)
                    
                    results = []
                    q = np.array(query_embedding, dtype=float)
                    
                    for emb_doc in emb_docs:
                        emb = emb_doc.get("embedding", [])
                        if not emb:
                            continue
                        
                        v = np.array(emb, dtype=float)
                        if v.shape != q.shape:
                            continue
                        
                        # Cosine similarity
                        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.5:  # Threshold for relevance
                            results.append({
                                "memo_id": emb_doc.get("memo_id"),
                                "score": sim,
                            })
                    
                    # Sort by similarity score
                    results.sort(key=lambda r: r["score"], reverse=True)
                    memo_ids_from_vector_search = [r["memo_id"] for r in results[:20]]
                    logger.info(f"✅ Vector search found {len(memo_ids_from_vector_search)} memos")
                else:
                    logger.warning("⚠️ OPENAI_API_KEY not set, falling back to regex search")
                    # Fallback to regex if no API key
                    query["content"] = {"$regex": content_search, "$options": "i"}
            except Exception as e:
                logger.warning(f"⚠️ Vector search failed: {e}, falling back to regex")
                # Fallback to regex on error
                query["content"] = {"$regex": content_search, "$options": "i"}
        else:
            # No content search, use regex fallback if needed
            pass

        # 3. Tag Filter
        if tag:
            # Tags trong mongo lưu là array ["tag1", "tag2"]
            # Query exact match trong array
            query["payload.tags"] = tag

        # 4. Execute Query
        if memo_ids_from_vector_search:
            # Use vector search results + date filter
            # Convert string IDs to ObjectId
            object_ids: list[ObjectId] = []
            for mid in memo_ids_from_vector_search:
                try:
                    if ObjectId.is_valid(mid):
                        object_ids.append(ObjectId(mid))
                except Exception:  # pragma: no cover - defensive
                    continue

            if not object_ids:
                # Vector search said there were matches but all IDs were invalid.
                # To avoid returning a broad set of memos by date/tag only,
                # short‑circuit and return an empty result set.
                logger.warning(
                    "Vector search returned only invalid memo IDs. "
                    "Returning empty result instead of broad date/tag query."
                )
                empty_payload = {
                    "status": "success",
                    "query": {
                        "start_date": start_date,
                        "end_date": end_date or start_date,
                        "content_search": content_search,
                        "tag": tag,
                    },
                    "count": 0,
                    "memos": [],
                }
                return json.dumps(empty_payload, ensure_ascii=False)

            query["_id"] = {"$in": object_ids}
            logger.info("MongoDB Query (with vector search): %s", query)
        else:
            logger.info("MongoDB Query: %s", query)
        
        cursor = mongodb_client.memos.find(query).sort("created_at", -1).limit(20)
        docs = await cursor.to_list(length=20)

        # 5. Format Result
        memos = []
        for doc in docs:
            memos.append(serialize_doc(doc))

        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,
        )

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