"""
Brain Client v2 - Connects Manus to the persistent memory server.
RAG Tier: Advanced (query expansion + re-ranking) + Hybrid (keyword fallback) + Agentic (iterative loop)
"""

import requests
import json
from datetime import datetime
from typing import Optional, List, Dict

# --- Server Configuration ---
POCKETBASE_URL = "https://cloud.mixpresso.top"
MEM0_URL = "https://mem.mixpresso.top"
BOOTSTRAP_URL = "bootstrap.mixpresso.top"
ADMIN_EMAIL = "admin@mixpresso.top"
ADMIN_PASS = "Mintyelmbrain"
USER_ID = "owner"
BRAIN_API_KEY = "7daf271ac08028dff270d69d4f8c8e3fc61821163a4c3dd5"
QDRANT_API_KEY = "b56a1e3a1671b4d0ed56ce5d44bfd318516aae8886dd045a8b4bde7daa6d02e7"

# --- RAG Configuration ---
RELEVANCE_THRESHOLD = 0.25   # below this triggers keyword fallback
HIGH_CONFIDENCE = 0.60       # above this = stop agentic loop early
MAX_AGENTIC_ITERATIONS = 3   # max search iterations in agentic mode

# --- Auth ---
_token_cache = {"token": None, "expires": 0}


def _get_token():
    """Get PocketBase auth token (cached)."""
    import time
    if _token_cache["token"] and time.time() < _token_cache["expires"]:
        return _token_cache["token"]

    r = requests.post(f"{POCKETBASE_URL}/api/collections/_superusers/auth-with-password", json={
        "identity": ADMIN_EMAIL, "password": ADMIN_PASS
    }, timeout=10)
    if r.status_code == 200:
        data = r.json()
        _token_cache["token"] = data["token"]
        _token_cache["expires"] = time.time() + 3500  # ~1 hour
        return data["token"]
    raise Exception(f"PocketBase auth failed: {r.status_code} {r.text[:200]}")


def _headers():
    return {"Authorization": _get_token(), "Content-Type": "application/json"}


# --- PocketBase CRUD ---

def save_record(collection: str, data: dict, record_id: str = None) -> dict:
    """Save or update a record in PocketBase."""
    data["last_updated"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    if record_id:
        r = requests.patch(f"{POCKETBASE_URL}/api/collections/{collection}/records/{record_id}",
                           headers=_headers(), json=data, timeout=10)
    else:
        r = requests.post(f"{POCKETBASE_URL}/api/collections/{collection}/records",
                          headers=_headers(), json=data, timeout=10)

    if r.status_code in (200, 201):
        return r.json()
    raise Exception(f"Save failed [{collection}]: {r.status_code} {r.text[:200]}")


def get_records(collection: str, filter_str: str = "", limit: int = 50) -> List[dict]:
    """Get records from PocketBase with optional filter."""
    params = {"perPage": limit}
    if filter_str:
        params["filter"] = filter_str
    r = requests.get(f"{POCKETBASE_URL}/api/collections/{collection}/records",
                     headers=_headers(), params=params, timeout=10)
    if r.status_code == 200:
        return r.json().get("items", [])
    return []


def get_record(collection: str, record_id: str) -> Optional[dict]:
    """Get a single record by ID."""
    r = requests.get(f"{POCKETBASE_URL}/api/collections/{collection}/records/{record_id}",
                     headers=_headers(), timeout=10)
    if r.status_code == 200:
        return r.json()
    return None


def delete_record(collection: str, record_id: str) -> bool:
    """Delete a record."""
    r = requests.delete(f"{POCKETBASE_URL}/api/collections/{collection}/records/{record_id}",
                        headers=_headers(), timeout=10)
    return r.status_code == 204


def find_record(collection: str, field: str, value: str) -> Optional[dict]:
    """Find a record by field value."""
    records = get_records(collection, filter_str=f'{field}="{value}"', limit=1)
    return records[0] if records else None


# --- Mem0 Semantic Search (Naive / Base Layer) ---

def index_memory(content: str, collection: str, chat_id: str = None,
                 category: str = None, tags: str = None, record_id: str = None) -> dict:
    """Index content into Qdrant for semantic search."""
    payload = {
        "content": content,
        "user_id": USER_ID,
        "collection": collection,
    }
    if chat_id:
        payload["chat_id"] = chat_id
    if category:
        payload["category"] = category
    if tags:
        payload["tags"] = tags
    if record_id:
        payload["record_id"] = record_id

    r = requests.post(f"{MEM0_URL}/memories/add", json=payload, timeout=15)
    if r.status_code == 200:
        return r.json()
    raise Exception(f"Index failed: {r.status_code} {r.text[:200]}")


def search_memory(query: str, collection: str = None, chat_id: str = None,
                  category: str = None, limit: int = 10) -> List[dict]:
    """Naive semantic search — single query, raw cosine similarity results."""
    payload = {
        "query": query,
        "user_id": USER_ID,
        "limit": limit,
    }
    if collection:
        payload["collection"] = collection
    if chat_id:
        payload["chat_id"] = chat_id
    if category:
        payload["category"] = category

    r = requests.post(f"{MEM0_URL}/memories/search", json=payload, timeout=15)
    if r.status_code == 200:
        return r.json().get("results", [])
    return []


def get_stats() -> dict:
    """Get memory statistics."""
    r = requests.get(f"{MEM0_URL}/stats", timeout=5)
    if r.status_code == 200:
        return r.json()
    return {}


# --- Advanced RAG: Query Expansion + Re-Ranking ---

def _expand_query(query: str) -> List[str]:
    """Generate semantic variants of the query using LLM. Falls back to original on error."""
    try:
        import os
        from openai import OpenAI
        client = OpenAI()  # uses OPENAI_API_KEY + OPENAI_API_BASE from env
        prompt = (
            f"Generate 3 different phrasings of this search query for a personal memory system.\n"
            f"Return only the 3 queries, one per line, no numbering, no explanation.\n"
            f"Original: {query}"
        )
        resp = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=150,
            temperature=0.3,
            timeout=5,
        )
        variants = resp.choices[0].message.content.strip().split("\n")
        variants = [v.strip() for v in variants if v.strip()][:3]
        return [query] + variants  # original + up to 3 variants
    except Exception:
        return [query]  # silent fallback to original query only


def advanced_search(query: str, collection: str = None, limit: int = 10) -> List[dict]:
    """
    Advanced RAG search: query expansion + multi-variant search + dedup + re-rank.
    Falls back to naive search if expansion fails.
    """
    # Step 1: Expand query into variants
    variants = _expand_query(query)

    # Step 2: Search all variants, collect and deduplicate results
    all_results: Dict[str, dict] = {}
    for i, variant in enumerate(variants):
        try:
            results = search_memory(variant, collection=collection, limit=limit)
        except Exception:
            continue
        for r in results:
            mid = r.get("memory_id", "")
            if not mid:
                continue
            if mid not in all_results:
                all_results[mid] = dict(r)
                all_results[mid]["_scores"] = []
            # Original query weighted 1.0x, variants 0.85x
            weight = 1.0 if i == 0 else 0.85
            all_results[mid]["_scores"].append(r.get("score", 0) * weight)

    if not all_results:
        return []

    # Step 3: Re-rank by best score across all variants
    ranked = sorted(
        all_results.values(),
        key=lambda x: max(x["_scores"]),
        reverse=True
    )

    # Step 4: Finalize scores, remove internal field
    for r in ranked:
        r["score"] = round(max(r.pop("_scores")), 4)

    return ranked[:limit]


# --- Hybrid RAG: Keyword Fallback via PocketBase ---

# Fields to search per collection for keyword matching
_KEYWORD_FIELDS = {
    "identity": ["content"],
    "facts": ["fact", "tags"],
    "area_tech": ["content", "topic"],
    "area_business": ["content", "topic"],
    "area_people": ["content", "topic"],
    "area_personal": ["content", "topic"],
    "projects": ["purpose", "tech_stack", "current_state", "conventions", "chat_name"],
    "episodes": ["summary", "problems_solved"],
    "decisions": ["decision", "reasoning"],
    "ideas": ["title", "description"],
}

_STOP_WORDS = {
    "what", "when", "where", "which", "that", "this", "with", "from",
    "have", "been", "were", "they", "their", "about", "into", "over",
    "after", "before", "more", "some", "such", "than", "then", "also",
    "each", "most", "other", "same", "just", "like", "will", "your",
}


def keyword_search(query: str, collection: str = None, limit: int = 10) -> List[dict]:
    """
    Keyword fallback search via PocketBase LIKE matching (~).
    Returns results in mem0-compatible format with fixed score of 0.5.
    """
    # Extract meaningful terms
    terms = [
        w.lower() for w in query.split()
        if len(w) > 3 and w.lower() not in _STOP_WORDS
    ]
    if not terms:
        return []

    collections_to_search = [collection] if collection else list(_KEYWORD_FIELDS.keys())
    seen_ids: set = set()
    results: List[dict] = []

    for coll in collections_to_search:
        fields = _KEYWORD_FIELDS.get(coll, ["content"])
        for term in terms[:3]:  # cap at 3 terms to keep queries simple
            filter_parts = [f'{field}~"{term}"' for field in fields]
            filter_str = " || ".join(filter_parts)
            try:
                records = get_records(coll, filter_str=filter_str, limit=limit)
            except Exception:
                continue
            for rec in records:
                rid = rec.get("id", "")
                if rid in seen_ids:
                    continue
                seen_ids.add(rid)
                # Use the primary text field as content
                primary_field = fields[0]
                content = rec.get(primary_field, "")
                results.append({
                    "memory_id": rid,
                    "content": content,
                    "collection": coll,
                    "score": 0.5,
                    "source": "keyword",
                    "record_id": rid,
                })

    return results[:limit]


def hybrid_search(query: str, collection: str = None, limit: int = 10) -> List[dict]:
    """
    Hybrid RAG: Advanced semantic search + keyword fallback when confidence is low.
    Merges both result sets, deduplicates, returns ranked results.
    """
    # Run advanced semantic search first
    semantic_results = advanced_search(query, collection=collection, limit=limit)

    # Check top score confidence
    top_score = semantic_results[0]["score"] if semantic_results else 0

    if top_score < RELEVANCE_THRESHOLD:
        # Low confidence — augment with keyword search
        kw_results = keyword_search(query, collection=collection, limit=limit)
        semantic_ids = {r["memory_id"] for r in semantic_results}
        for kr in kw_results:
            if kr["memory_id"] not in semantic_ids:
                semantic_results.append(kr)

    return semantic_results[:limit]


# --- Agentic RAG: Iterative Search with Quality Evaluation ---

def _evaluate_results(query: str, results: List[dict]) -> dict:
    """
    Use LLM to evaluate if retrieved results actually answer the query.
    Returns {"quality": "good"|"poor", "refined_query": str}
    Falls back to {"quality": "good"} on any error.
    """
    if not results:
        return {"quality": "poor", "refined_query": query}
    try:
        from openai import OpenAI
        client = OpenAI()
        context = "\n".join([f"- {r.get('content', '')[:150]}" for r in results[:5]])
        prompt = (
            f"You are evaluating memory retrieval results for a personal AI assistant.\n\n"
            f"Query: {query}\n\n"
            f"Retrieved results:\n{context}\n\n"
            f"Are these results actually relevant to the query?\n"
            f"If not, what refined query would find better results?\n\n"
            f'Respond in JSON only: {{"relevant": true or false, "refined_query": "..."}}'
        )
        resp = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=100,
            temperature=0,
            response_format={"type": "json_object"},
            timeout=5,
        )
        data = json.loads(resp.choices[0].message.content)
        quality = "good" if data.get("relevant", True) else "poor"
        refined = data.get("refined_query", query) or query
        return {"quality": quality, "refined_query": refined}
    except Exception:
        return {"quality": "good", "refined_query": query}


def agentic_search(query: str, collection: str = None, limit: int = 10) -> List[dict]:
    """
    Agentic RAG: Iterative search with LLM quality evaluation and query refinement.
    Runs up to MAX_AGENTIC_ITERATIONS iterations, exits early on high confidence.
    Falls back to hybrid_search results if iteration produces no improvement.
    """
    best_results: List[dict] = []
    best_score: float = 0.0
    current_query = query
    queries_tried: set = set()

    for iteration in range(MAX_AGENTIC_ITERATIONS):
        if current_query in queries_tried:
            break
        queries_tried.add(current_query)

        # Use hybrid search for each iteration (gets both semantic + keyword)
        results = hybrid_search(current_query, collection=collection, limit=limit)

        # Track best results across iterations
        top_score = results[0]["score"] if results else 0.0
        if top_score > best_score:
            best_score = top_score
            best_results = results

        # Early exit if high confidence
        if top_score >= HIGH_CONFIDENCE:
            break

        # Evaluate quality and get refined query for next iteration
        evaluation = _evaluate_results(current_query, results)
        if evaluation["quality"] == "good":
            break

        refined = evaluation.get("refined_query", current_query)
        if refined == current_query:
            break  # No new query to try
        current_query = refined

    return best_results


# --- Public Search Interface ---

def brain_search(query: str, collection: str = None, limit: int = 10,
                 mode: str = "advanced") -> List[dict]:
    """
    Search the brain.

    Modes:
      "naive"    — single raw semantic search (fastest, least accurate)
      "advanced" — query expansion + re-ranking (default, ~600ms)
      "hybrid"   — advanced + keyword fallback when confidence is low
      "agentic"  — iterative search with quality evaluation (~1-2s, most thorough)

    All modes degrade gracefully to naive on failure.
    """
    try:
        if mode == "agentic":
            return agentic_search(query, collection=collection, limit=limit)
        elif mode == "hybrid":
            return hybrid_search(query, collection=collection, limit=limit)
        elif mode == "advanced":
            return advanced_search(query, collection=collection, limit=limit)
        else:
            return search_memory(query, collection=collection, limit=limit)
    except Exception:
        # Last-resort fallback to naive search
        return search_memory(query, collection=collection, limit=limit)


# --- High-Level Memory Operations ---

def save_to_identity(section: str, content: str, tags: str = "") -> dict:
    """Save to the master brain identity (always loads)."""
    existing = find_record("identity", "section", section)
    record = save_record("identity", {"section": section, "content": content, "tags": tags},
                         record_id=existing["id"] if existing else None)
    index_memory(f"[Identity - {section}] {content}", "identity",
                 category=section, tags=tags, record_id=record["id"])
    return record


def save_secret(name: str, value: str, category: str = "other",
                related_service: str = "", notes: str = "") -> dict:
    """Save a credential/secret."""
    existing = find_record("secrets_vault", "name", name)
    record = save_record("secrets_vault", {
        "name": name, "value": value, "category": category,
        "related_service": related_service, "notes": notes, "status": "active"
    }, record_id=existing["id"] if existing else None)
    index_memory(f"[Secret] {name} for {related_service}. {notes}",
                 "secrets_vault", category=category, record_id=record["id"])
    return record


def save_project_memory(chat_id: str, chat_name: str, **fields) -> dict:
    """Save/update project memory for a specific chat."""
    existing = find_record("projects", "chat_id", chat_id)
    data = {"chat_id": chat_id, "chat_name": chat_name}
    data.update(fields)
    record = save_record("projects", data,
                         record_id=existing["id"] if existing else None)
    searchable = f"[Project: {chat_name}] "
    for key, val in fields.items():
        if val:
            searchable += f"{key}: {val}. "
    index_memory(searchable, "projects", chat_id=chat_id, record_id=record["id"])
    return record


def save_fact(fact: str, category: str = "technical", source_chat: str = "",
              confidence: str = "confirmed", tags: str = "") -> dict:
    """Save a learned fact."""
    record = save_record("facts", {
        "fact": fact, "category": category, "source_chat": source_chat,
        "confidence": confidence, "tags": tags, "learned_on": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    })
    index_memory(f"[Fact - {category}] {fact}", "facts",
                 category=category, tags=tags, record_id=record["id"])
    return record


def save_episode(chat_id: str, chat_name: str, summary: str,
                 problems_solved: str = "", tags: str = "") -> dict:
    """Save a session episode."""
    record = save_record("episodes", {
        "chat_id": chat_id, "chat_name": chat_name,
        "session_date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "summary": summary, "problems_solved": problems_solved, "tags": tags
    })
    index_memory(f"[Episode - {chat_name}] {summary}", "episodes",
                 chat_id=chat_id, tags=tags, record_id=record["id"])
    return record


def save_decision(decision: str, reasoning: str = "", alternatives: str = "",
                  source_chat: str = "", tags: str = "") -> dict:
    """Save a key decision."""
    record = save_record("decisions", {
        "decision": decision, "reasoning": reasoning,
        "alternatives_considered": alternatives, "source_chat": source_chat,
        "date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "status": "active", "tags": tags
    })
    index_memory(f"[Decision] {decision}. Reason: {reasoning}", "decisions",
                 tags=tags, record_id=record["id"])
    return record


def save_idea(title: str, description: str = "", category: str = "project",
              tags: str = "") -> dict:
    """Save an idea or brainstorm."""
    record = save_record("ideas", {
        "title": title, "description": description, "category": category,
        "status": "new", "tags": tags, "created": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    })
    index_memory(f"[Idea - {category}] {title}: {description}", "ideas",
                 category=category, tags=tags, record_id=record["id"])
    return record


def save_area(area: str, topic: str, content: str, category: str = "",
              tags: str = "") -> dict:
    """Save to a life area (area_tech, area_business, area_people, area_personal)."""
    collection = f"area_{area}"
    existing = find_record(collection, "topic", topic)
    data = {"topic": topic, "content": content, "tags": tags}
    if category:
        data["category"] = category
    record = save_record(collection, data,
                         record_id=existing["id"] if existing else None)
    index_memory(f"[{area.title()} - {topic}] {content}", collection,
                 category=category, tags=tags, record_id=record["id"])
    return record


# --- Load Operations ---

def load_identity() -> List[dict]:
    """Load the full identity/master brain."""
    return get_records("identity", limit=50)


def load_secrets() -> List[dict]:
    """Load all active secrets."""
    return get_records("secrets_vault", filter_str='status="active"', limit=100)


def load_project(chat_id: str) -> Optional[dict]:
    """Load project memory for a specific chat."""
    return find_record("projects", "chat_id", chat_id)


def load_area(area: str) -> List[dict]:
    """Load all records from a life area."""
    return get_records(f"area_{area}", limit=50)


# --- Status ---

def status() -> dict:
    """Get full brain status."""
    stats = get_stats()
    collections_status = {}
    for coll in ["identity", "secrets_vault", "area_tech", "area_business",
                 "area_people", "area_personal", "projects", "facts",
                 "episodes", "decisions", "ideas", "archive"]:
        collections_status[coll] = len(get_records(coll, limit=200))

    return {
        "server": "connected",
        "pocketbase": "healthy",
        "mem0": stats.get("status", "unknown"),
        "total_vectors": stats.get("total_memories", 0),
        "rag_mode": "Advanced + Hybrid + Agentic",
        "collections": collections_status,
    }