"""Memory-aware extraction — informed by existing knowledge.

Enriches fact extraction with existing entity context so the LLM can
avoid re-extracting known information and produce more precise facts.
Reconciliation (ADD/UPDATE/NOOP/DELETE) is handled by ``reconcile_facts``
downstream — this module only extracts.

Pipeline: alias lookup (DB) → pre-retrieval (pgvector) → informed LLM call.
Falls back to blind extraction on any failure (caller handles fallback).
"""

import json
import logging
import re
from datetime import UTC, datetime

from sqlalchemy import select, text
from sqlalchemy.ext.asyncio import AsyncSession

from arandu._helpers import strip_markdown_fences
from arandu.config import MemoryConfig
from arandu.constants import (
    ExtractedEntity,
    ExtractedFact,
    ExtractedRelation,
    ExtractionOutput,
)
from arandu.models import MemoryEntity, MemoryEntityAlias
from arandu.protocols import EmbeddingProvider, LLMProvider
from arandu.write.extract import _call_llm, _deduplicate_facts, _normalize_extraction_subjects

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Token estimation
# ---------------------------------------------------------------------------

_AVG_CHARS_PER_TOKEN = 4


def _est_tokens(s: str) -> int:
    """Heuristic token count: ~4 chars per token."""
    return max(0, len(s) // _AVG_CHARS_PER_TOKEN)


# ---------------------------------------------------------------------------
# Prompt
# ---------------------------------------------------------------------------

INFORMED_EXTRACTION_PROMPT = """You are a personal knowledge extractor with access to existing memory.
Given a conversation message AND existing knowledge about known entities, extract ONLY genuinely new information.

SPEAKER CONTEXT: The speaker of this message is "{speaker_name}".
All first-person pronouns (I, me, my, myself, eu, mim, meu, minha, yo, mi, je, moi, ich) refer to "{speaker_name}".
Second-person pronouns (you, your, você, tu, teu, tua, seu, sua) refer to the INTERLOCUTOR (the person the speaker is talking TO), NOT to the speaker. Use "Recent conversation context" to identify who the interlocutor is. Do NOT attribute interlocutor's actions/plans to the speaker.

TEMPORAL CONTEXT: This message was sent on {message_date}.
When the message contains relative time references ("yesterday", "last week", "three years ago", "recently", etc.), resolve them to absolute dates based on the message date above. Include the resolved date in the fact text.

{existing_context}

RULES FOR ENTITIES:
- Identify ALL entities mentioned (person, place, organization, product, event, concept, pet, team).
- The speaker "{speaker_name}" MUST be included with type="person".
- NEVER create an entity called "user" — use "{speaker_name}" instead.
- NEVER create entities from pronouns — they refer to "{speaker_name}".
- When the same entity is referred to by multiple names, create ONE entity with the most complete name as "name" and alternatives in "aliases".

RULES FOR FACTS:
- Extract ALL factual information from the message, including specific details.
- The system will handle deduplication — your job is to CAPTURE everything, not filter.
- ONE fact = ONE piece of information. Never combine multiple pieces.
- Each fact has: subject, fact, confidence, importance_category.
- Use "{speaker_name}" as subject for facts about the speaker. NEVER use "user".
- SELF-CONTAINED: Fact text MUST include the subject's name.
- UNAMBIGUOUS: Include specific type, name, or qualifier.
- PRESERVE SPECIFIC DETAILS: Always keep proper nouns, titles (books, movies, songs, shows), country/city names, organization names, pet names, dates, numbers, and specific descriptions. "Melanie read Nothing is Impossible" is better than "Melanie read a book".
- Confidence: 0.95 (explicitly stated), 0.80 (strongly implied), 0.60 (weakly implied).
- NO DUPLICATE INFORMATION: each piece of info in exactly ONE fact, attributed to the ACTOR/POSSESSOR.
- NEVER create mirror facts from active and passive perspectives.

RULES FOR DEDUPLICATION:
- If information is already known exactly as-is in EXISTING KNOWLEDGE: do NOT extract it at all.
- The system will reconcile new vs existing facts — your job is to extract, not to decide ADD/UPDATE.

RULES FOR IMPORTANCE:
- "biographical_milestone": marriage, birth, death, graduation, first steps, major life event.
- "relationship_change": new friend, breakup, new job, joining a group.
- "stable_preference": likes art, interested in mental health, personality traits.
- "specific_event": went to parade, saw meteor shower, visited somewhere.
- "routine_activity": went to beach, had dinner, went for a walk.
- "conversational": greeting, thanking, agreeing, small talk.

RULES FOR RELATIONS:
- Extract EVERY relationship mentioned or strongly implied between entities.
- Use short, descriptive snake_case types (works_at, lives_in, friend_of, member_of, etc.).
- Use "{speaker_name}" not "user".
- If a relationship already exists in EXISTING KNOWLEDGE, do NOT re-extract it.

RULES FOR ENTITY PROFILES:
- For each entity you extract facts about, provide an updated profile in "updated_profiles".
- A profile is a CUMULATIVE summary (100-300 tokens) of who/what the entity is.
- IMPORTANT: When an existing profile is provided above, your updated profile MUST PRESERVE all information from the existing profile AND incorporate new facts from this message. Do NOT replace the existing profile — MERGE new information into it.
- If an entity's existing profile is provided and the new message adds nothing new, set that entity to null in updated_profiles.
- If an entity has no existing profile, create one covering all known facts about that entity.
- Be factual. Do NOT invent information not supported by the facts.
- Write in English.

Respond in JSON:
{{
  "entities": [{{"name": ..., "type": ..., "aliases": [...]}}],
  "facts": [{{"subject": ..., "fact": ..., "confidence": ..., "importance_category": ...}}],
  "relations": [{{"source": ..., "target": ..., "rel_type": ...}}],
  "updated_profiles": {{"entity_name": "updated profile text"|null}}
}}"""


# ---------------------------------------------------------------------------
# Alias lookup
# ---------------------------------------------------------------------------


async def alias_lookup(
    session: AsyncSession,
    agent_id: str,
    message: str,
) -> dict[str, str]:
    """Find known entities in a message via word-boundary matching on aliases.

    Loads all aliases for the user, skips those shorter than 3 characters,
    and checks for word-boundary matches in the message text.

    Args:
        session: Database session.
        agent_id: User identifier.
        message: Message text to scan.

    Returns:
        Dict mapping canonical_entity_key to the matched alias text.
    """
    stmt = select(
        MemoryEntityAlias.alias,
        MemoryEntityAlias.canonical_entity_key,
    ).where(MemoryEntityAlias.agent_id == agent_id)

    result = await session.execute(stmt)
    aliases = list(result.all())

    if not aliases:
        return {}

    message_lower = message.lower()
    matched: dict[str, str] = {}

    for alias_text, entity_key in aliases:
        if len(alias_text) < 3:
            continue

        alias_lower = alias_text.lower().strip()
        if not alias_lower:
            continue

        # Word boundary matching to avoid false positives
        pattern = r"\b" + re.escape(alias_lower) + r"\b"
        if re.search(pattern, message_lower):
            if entity_key not in matched or len(alias_text) > len(matched[entity_key]):
                matched[entity_key] = alias_text

    if matched:
        logger.info(
            "alias_lookup: found %d known entities: %s",
            len(matched),
            list(matched.keys()),
        )

    return matched


# ---------------------------------------------------------------------------
# Entity profile loader
# ---------------------------------------------------------------------------


async def load_entity_profiles(
    session: AsyncSession,
    agent_id: str,
    entity_keys: list[str],
) -> dict[str, str]:
    """Load existing profiles for known entities.

    Args:
        session: Database session.
        agent_id: User identifier.
        entity_keys: Canonical keys of entities to load profiles for.

    Returns:
        Dict mapping canonical_entity_key to profile_text (only non-null profiles).
    """
    if not entity_keys:
        return {}

    stmt = select(
        MemoryEntity.canonical_key,
        MemoryEntity.profile_text,
    ).where(
        MemoryEntity.agent_id == agent_id,
        MemoryEntity.canonical_key.in_(entity_keys),
        MemoryEntity.profile_text.isnot(None),
    )
    result = await session.execute(stmt)
    return {row[0]: row[1] for row in result.all()}


# ---------------------------------------------------------------------------
# Pre-retrieval context builder
# ---------------------------------------------------------------------------


async def build_extraction_context(
    session: AsyncSession,
    agent_id: str,
    entity_keys: list[str],
    message_embedding: list[float],
    config: MemoryConfig,
    entity_profiles: dict[str, str] | None = None,
) -> str:
    """Build existing knowledge context for the informed extraction prompt.

    For each entity, uses the profile_text if available (more compact and
    coherent). Falls back to top-K facts by embedding similarity when no
    profile exists.

    Args:
        session: Database session.
        agent_id: User identifier.
        entity_keys: Canonical keys of entities found in the message.
        message_embedding: Embedding vector of the incoming message.
        config: Memory configuration (topk, budget).
        entity_profiles: Pre-loaded profiles (entity_key → profile_text).
            Entities with profiles skip the fact retrieval.

    Returns:
        Formatted context string, or empty string if no relevant facts found.
    """
    if not entity_keys or not message_embedding:
        return ""

    profiles = entity_profiles or {}
    topk = config.informed_extraction_topk
    budget = config.informed_extraction_context_budget_tokens

    per_entity_limit = topk
    if len(entity_keys) > 2:
        per_entity_limit = max(3, budget // len(entity_keys) // 20)

    context_parts: list[str] = []
    total_tokens = 0

    for entity_key in entity_keys:
        if total_tokens >= budget:
            break

        entity_display = entity_key.split(":")[-1].replace("_", " ").title()

        # Use profile if available (more compact and coherent)
        if entity_key in profiles:
            profile = profiles[entity_key]
            block = f"Profile of {entity_display}:\n  {profile}"
            block_tokens = _est_tokens(block)
            if total_tokens + block_tokens <= budget:
                context_parts.append(block)
                total_tokens += block_tokens
            continue

        # Fallback: top-K facts by embedding similarity
        query = text("""
            SELECT fact_text,
                   1 - (embedding_vec <=> CAST(:embedding AS vector)) as similarity
            FROM memory_facts
            WHERE agent_id = :agent_id
              AND entity_key = :entity_key
              AND valid_to IS NULL
              AND embedding_vec IS NOT NULL
            ORDER BY embedding_vec <=> CAST(:embedding AS vector)
            LIMIT :limit
        """)

        result = await session.execute(
            query,
            {
                "agent_id": agent_id,
                "entity_key": entity_key,
                "embedding": str(message_embedding),
                "limit": per_entity_limit,
            },
        )
        rows = list(result.all())

        if not rows:
            continue

        facts = [(row[0], row[1]) for row in rows if row[1] >= 0.3]
        if not facts:
            continue

        lines = [f"  - {fact_text}" for fact_text, _ in facts]
        block = f"Known facts about {entity_display}:\n" + "\n".join(lines)

        block_tokens = _est_tokens(block)
        if total_tokens + block_tokens > budget:
            remaining = budget - total_tokens
            truncated_lines: list[str] = []
            for line in lines:
                line_tokens = _est_tokens(line)
                if remaining - line_tokens < 0:
                    break
                truncated_lines.append(line)
                remaining -= line_tokens
            if truncated_lines:
                block = f"Known facts about {entity_display}:\n" + "\n".join(truncated_lines)
                block_tokens = _est_tokens(block)
            else:
                continue

        context_parts.append(block)
        total_tokens += block_tokens

    if not context_parts:
        return ""

    return "EXISTING KNOWLEDGE (do NOT re-extract what is already known):\n\n" + "\n\n".join(context_parts)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


async def informed_extract_from_message(
    session: AsyncSession,
    agent_id: str,
    message: str,
    llm: LLMProvider,
    embeddings: EmbeddingProvider,
    speaker_name: str,
    config: MemoryConfig,
    recent_messages: list[str] | None = None,
    occurred_at: datetime | None = None,
) -> tuple[ExtractionOutput, dict]:
    """Extract entities, facts, and relations informed by existing memory.

    Pipeline:
        1. Alias lookup — find known entities in the message (no LLM).
        2. Pre-retrieval — fetch relevant existing facts for context (pgvector).
        3. Informed extraction — single LLM call with existing context.

    Args:
        session: Database session.
        agent_id: User identifier.
        message: The user message to extract from.
        llm: Injected LLM provider.
        embeddings: Injected embedding provider.
        speaker_name: Name of the speaker.
        config: Memory configuration.
        recent_messages: Optional conversation context (last N messages).

    Returns:
        Tuple of (ExtractionOutput, extraction_meta dict).

    Raises:
        Exception on any failure — caller should catch and fall back to
        blind extraction.
    """
    # 1. Alias lookup
    matched_entities = await alias_lookup(session, agent_id, message)

    # 2. Load profiles + build context
    existing_context = ""
    entity_profiles: dict[str, str] = {}
    if matched_entities:
        entity_profiles = await load_entity_profiles(session, agent_id, list(matched_entities.keys()))
        message_embedding = await embeddings.embed_one(message)
        if message_embedding:
            existing_context = await build_extraction_context(
                session,
                agent_id,
                list(matched_entities.keys()),
                message_embedding,
                config,
                entity_profiles=entity_profiles,
            )

    # 3. Build prompt
    if existing_context:
        context_block = existing_context
    else:
        context_block = "EXISTING KNOWLEDGE: None (first interaction or no known entities detected)."

    conv_context = ""
    if recent_messages:
        trimmed = [m[:300] for m in recent_messages[-5:]]
        conv_context = "\n\nRecent conversation context:\n" + "\n".join(f"- {m}" for m in trimmed)

    message_date = (occurred_at or datetime.now(UTC)).strftime("%B %d, %Y")
    system_prompt = INFORMED_EXTRACTION_PROMPT.format(
        speaker_name=speaker_name,
        existing_context=context_block,
        message_date=message_date,
    )
    user_prompt = f"Current message to extract from:\n{message}{conv_context}"

    # 4. Call LLM
    timeout = config.extraction_timeout_sec if config.extraction_timeout_sec > 0 else None
    raw = await _call_llm(
        llm=llm,
        system_prompt=system_prompt,
        user_prompt=user_prompt,
        call_type="informed_extraction",
        timeout_sec=timeout,
    )

    if raw is None:
        raise ValueError("Informed extraction LLM call returned None")

    # 5. Parse output
    cleaned = strip_markdown_fences(raw)
    data = json.loads(cleaned)

    entities = [ExtractedEntity.model_validate(e) for e in (data.get("entities") or [])]
    facts = [ExtractedFact.model_validate(f) for f in (data.get("facts") or [])]
    relations = [ExtractedRelation.model_validate(r) for r in (data.get("relations") or [])]

    output = ExtractionOutput(entities=entities, facts=facts, relations=relations)

    # 6. Parse updated_profiles (optional, fail-safe)
    updated_profiles: dict[str, str | None] = {}
    raw_profiles = data.get("updated_profiles")
    if isinstance(raw_profiles, dict):
        for name, profile in raw_profiles.items():
            if profile is None or isinstance(profile, str):
                updated_profiles[name] = profile

    # 7. Normalize subjects + deduplicate
    output = _normalize_extraction_subjects(output)
    output.facts = _deduplicate_facts(list(output.facts))

    logger.info(
        "informed_extraction: entities=%d facts=%d relations=%d context_entities=%d profiles_updated=%d",
        len(output.entities),
        len(output.facts),
        len(output.relations),
        len(matched_entities),
        sum(1 for v in updated_profiles.values() if v is not None),
    )

    extraction_meta: dict = {
        "mode": "informed",
        "context_entities": len(matched_entities),
        "context_entity_keys": list(matched_entities.keys()),
        "updated_profiles": updated_profiles,
    }

    return output, extraction_meta
