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

Replaces blind extraction + reconciliation with a single LLM call that
receives existing entity context and extracts only the delta (genuinely
new or updated information). Classifies importance at extraction time.

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
import uuid as uuid_mod

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 (
    IMPORTANCE_CATEGORY_DEFAULT,
    IMPORTANCE_CATEGORY_MAP,
    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
from arandu.write.reconcile import ReconciliationDecision

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}".

{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 ONLY genuinely new information NOT already covered by existing knowledge above.
- If information is already known, do NOT extract it.
- ONE fact = ONE piece of information. Never combine multiple pieces.
- Each fact has: subject, fact, confidence, action, 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.
- 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 ACTION:
- "NEW": genuinely new information not in existing knowledge.
- "UPDATE": replaces or refines an existing fact. Set "updates" to the existing fact text being replaced.
- If information is already known exactly as-is: do NOT extract it at all.

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 concise summary (100-300 tokens) of who/what the entity is, based on ALL known information (existing profile + new facts).
- If an entity's existing profile is provided above and the new message does NOT change the understanding, set that entity to null in updated_profiles.
- If an entity has no existing profile, create one from the extracted facts.
- Be factual. Do NOT invent information not supported by the facts.
- Write in the same language as the conversation.

Respond in JSON:
{{
  "entities": [{{"name": ..., "type": ..., "aliases": [...]}}],
  "facts": [{{"subject": ..., "fact": ..., "confidence": ..., "action": "NEW"|"UPDATE", "importance_category": ..., "updates": null|"existing fact text"}}],
  "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)


# ---------------------------------------------------------------------------
# Decision builder (informed facts → ReconciliationDecision)
# ---------------------------------------------------------------------------


async def build_informed_decisions(
    session: AsyncSession,
    agent_id: str,
    facts: list[ExtractedFact],
    entity_map: dict,
    embeddings: EmbeddingProvider,
) -> list[ReconciliationDecision]:
    """Convert informed extraction facts to ReconciliationDecision format.

    NEW facts become ADD decisions. UPDATE facts search for the existing
    fact by embedding similarity and become UPDATE decisions (with
    matched_fact_id) or fall back to ADD if the target is not found.

    Args:
        session: Database session.
        agent_id: User identifier.
        facts: Extracted facts with action and importance_category.
        entity_map: Resolved entity map (name → ResolvedEntity).
        embeddings: Embedding provider for UPDATE target search.

    Returns:
        List of ReconciliationDecision ready for execute_upsert.
    """
    decisions: list[ReconciliationDecision] = []

    for fact in facts:
        importance = IMPORTANCE_CATEGORY_MAP.get(fact.importance_category, IMPORTANCE_CATEGORY_DEFAULT)

        # Resolve subject to entity key
        resolved = entity_map.get(fact.subject)
        subject_key = resolved.entity_key if resolved else fact.subject

        if fact.action == "UPDATE" and fact.updates:
            matched_id = await _find_update_target(session, agent_id, subject_key, fact.updates, embeddings)
            if matched_id:
                decisions.append(
                    ReconciliationDecision(
                        fact_text=fact.fact,
                        subject=subject_key,
                        action="UPDATE",
                        matched_fact_id=matched_id,
                        matched_fact_text=fact.updates,
                        confidence=fact.confidence,
                        importance=importance,
                    )
                )
            else:
                # Target not found — fall back to ADD
                decisions.append(
                    ReconciliationDecision(
                        fact_text=fact.fact,
                        subject=subject_key,
                        action="ADD",
                        confidence=fact.confidence,
                        importance=importance,
                    )
                )
        else:
            # NEW → ADD
            decisions.append(
                ReconciliationDecision(
                    fact_text=fact.fact,
                    subject=subject_key,
                    action="ADD",
                    confidence=fact.confidence,
                    importance=importance,
                )
            )

    return decisions


async def _find_update_target(
    session: AsyncSession,
    agent_id: str,
    entity_key: str,
    updates_text: str,
    embeddings: EmbeddingProvider,
) -> uuid_mod.UUID | None:
    """Find the existing fact that an UPDATE action refers to.

    Embeds the updates text and searches for the most similar active fact
    belonging to the entity. Returns the fact ID if similarity >= 0.5.
    """
    embedding = await embeddings.embed_one(updates_text)
    if not embedding:
        return None

    query = text("""
        SELECT id,
               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 1
    """)

    result = await session.execute(
        query,
        {
            "agent_id": agent_id,
            "entity_key": entity_key,
            "embedding": str(embedding),
        },
    )
    row = result.first()
    if row and row[1] >= 0.5:
        return uuid_mod.UUID(str(row[0]))
    return None


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

    system_prompt = INFORMED_EXTRACTION_PROMPT.format(
        speaker_name=speaker_name,
        existing_context=context_block,
    )
    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
