"""Memory extraction — entities, facts, and relations from a message.

Pipeline: Entity scan → fact extraction → relation extraction (concurrent).
All entities are processed in a single fact extraction call (no batching).

Ported from advisor_api extract.py. All LLM calls go through the injected
LLMProvider — no global imports.
"""

import asyncio
import json
import logging

from arandu._helpers import strip_markdown_fences
from arandu.config import MemoryConfig
from arandu.constants import (
    ExtractedEntity,
    ExtractedFact,
    ExtractedRelation,
    ExtractionOutput,
)
from arandu.protocols import LLMProvider

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Prompts — single-pass
# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
# Prompts
# ---------------------------------------------------------------------------
ENTITY_SCAN_PROMPT = """You are a personal knowledge extractor.
Given a conversation message, identify ALL entities mentioned.

ENTITIES: Every person, place, organization, product, event, concept, pet, team, or named thing.
- Each entity has: name, type (free-form string, e.g. person, organization, place, product, event, concept, pet), aliases (optional list of alternative names)
- The user themselves is always included: name="user", type="person"
- Include EVERY named entity, even if mentioned only once.
- IMPORTANT: When the same entity is referred to by multiple names in the message (nickname + full name, abbreviation + full form, colloquial + formal, role + name), create ONE entity with the most complete/formal name as "name" and the other names in "aliases".
  Examples:
  - "my friend Guili (Guilherme Maturana)" → {"name": "Guilherme Maturana", "type": "person", "aliases": ["Guili"]}
  - "my therapist Dr. Carla" and later "Carla said..." → {"name": "Dr. Carla", "type": "person", "aliases": ["Carla"]}
  - "Stone (StoneCo)" → {"name": "StoneCo", "type": "organization", "aliases": ["Stone"]}
- Only group names when the message explicitly indicates they refer to the same entity. Do not guess.
- Pay careful attention to entity type classification. Cities, countries, and neighborhoods are "place", not "person" — even when mentioned alongside many people.

Respond in JSON with key: entities (array of {name, type, aliases})."""

FACT_EXTRACTION_PROMPT = """You are a personal knowledge extractor.
Given a conversation message and a list of entities, extract ALL atomic facts.

Rules:
- ONE fact = ONE piece of information. Never combine multiple pieces into one fact.
- Each fact has: subject (entity name), fact (natural language sentence), confidence (0.0-1.0)
- SELF-CONTAINED: The fact text MUST include the subject's name. "Patrícia Duarte is a cardiologist", NOT "is a cardiologist".
- Use "user" as subject for facts about the user themselves.
- Confidence: 0.95 (explicitly stated), 0.80 (strongly implied), 0.60 (weakly implied).
- Every relationship between entities MUST also generate a corresponding fact from the source entity's perspective. "Sarah is my wife" → fact: "Sarah is user's wife".

NO DUPLICATE INFORMATION: Each piece of information MUST appear in exactly ONE fact.
- "Clara left Vertix" is ONE fact with subject "Clara Rezende". Do NOT also create "Vertix lost Clara Rezende" — that is the same information rephrased.
- If entity A did something involving entity B, the fact belongs to the ACTOR. "Thiago hired Clara" → subject is Thiago. Do NOT also create "Clara was hired by Thiago".
- NEVER extract the same information from both the active and passive perspective. Pick ONE.
- When in doubt, use the grammatical subject of the original sentence.

Respond in JSON with key: facts (array of {subject, fact, confidence})."""

RELATION_EXTRACTION_PROMPT = """You are a personal knowledge extractor.
Given a conversation message and a list of entities, extract ALL relationships between them.

Each relation has: source, target, rel_type
- Use a short, descriptive snake_case type (e.g., works_at, lives_in, spouse_of, member_of, friend_of, owns).
- Use any type that accurately describes the relationship. Be specific.

Rules:
- Extract EVERY relationship mentioned or strongly implied.
- "A works at B" → source="A", target="B", rel_type="works_at"
- "A is part of B" → source="A", target="B", rel_type="member_of"
- If the message implies a chain (A → B → C), extract each link separately.

Respond in JSON with key: relations (array of {source, target, rel_type})."""


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


# Internal alias — used within this module as _strip_markdown_fences.
_strip_markdown_fences = strip_markdown_fences


def _parse_extraction_output(raw: str) -> ExtractionOutput:
    """Parse LLM JSON response into ExtractionOutput."""
    cleaned = _strip_markdown_fences(raw)
    data = json.loads(cleaned)
    return ExtractionOutput.model_validate(data)


# Relation types that indicate identity (alias/same entity).
_IDENTITY_REL_TYPES = frozenset(
    {
        "same_as",
        "also_known_as",
        "nickname_of",
        "alias_of",
    }
)


def _normalize_extraction_subjects(output: ExtractionOutput) -> ExtractionOutput:
    """Remap fact subjects and relation sources/targets from aliases to canonical entity names.

    Also removes identity relations (same_as, also_known_as, etc.) where both
    source and target resolve to the same canonical entity after normalization.

    Args:
        output: Raw extraction output (may contain aliases as subjects).

    Returns:
        Normalized ExtractionOutput with canonical names in subjects/sources/targets.
    """
    # Build alias → canonical name map (case-insensitive)
    alias_map: dict[str, str] = {}
    for entity in output.entities:
        for alias in entity.aliases:
            alias_map[alias.lower().strip()] = entity.name

    if not alias_map:
        return output

    def _resolve(name: str) -> str:
        return alias_map.get(name.lower().strip(), name)

    # Remap fact subjects
    normalized_facts = [
        ExtractedFact(subject=_resolve(f.subject), fact=f.fact, confidence=f.confidence) for f in output.facts
    ]

    # Remap relation sources/targets, remove self-referencing identity relations
    normalized_relations = []
    for rel in output.relations:
        resolved_source = _resolve(rel.source)
        resolved_target = _resolve(rel.target)
        if (
            rel.rel_type.lower().strip() in _IDENTITY_REL_TYPES
            and resolved_source.lower().strip() == resolved_target.lower().strip()
        ):
            continue
        normalized_relations.append(
            ExtractedRelation(source=resolved_source, target=resolved_target, rel_type=rel.rel_type)
        )

    return ExtractionOutput(
        entities=output.entities,
        facts=normalized_facts,
        relations=normalized_relations,
    )


def _deduplicate_facts(facts: list[ExtractedFact]) -> list[ExtractedFact]:
    """Remove duplicate facts by normalized (subject, fact_text) key.

    Normalization: lowercase, strip whitespace and trailing punctuation.
    Keeps first occurrence, discards subsequent duplicates.
    """
    seen: set[tuple[str, str]] = set()
    deduped: list[ExtractedFact] = []
    for f in facts:
        key = (
            f.subject.lower().strip().rstrip(".,;:!?"),
            f.fact.lower().strip().rstrip(".,;:!?"),
        )
        if key not in seen:
            seen.add(key)
            deduped.append(f)
    return deduped


async def _call_llm(
    llm: LLMProvider,
    system_prompt: str,
    user_prompt: str,
    call_type: str = "extraction",
    timeout_sec: float | None = None,
) -> str | None:
    """Call LLM with JSON mode. Returns raw content string or None on failure."""
    try:
        coro = llm.complete(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0,
            response_format={"type": "json_object"},
        )
        if timeout_sec and timeout_sec > 0:
            result = await asyncio.wait_for(coro, timeout=timeout_sec)
        else:
            result = await coro
        return result.text
    except TimeoutError:
        logger.warning("extraction: LLM call timed out after %.1fs (%s)", timeout_sec, call_type)
        return None
    except Exception as e:
        logger.warning("extraction: LLM call failed (%s): %s", call_type, e)
        return None


# ---------------------------------------------------------------------------
# Multi-pass extraction pipeline
# ---------------------------------------------------------------------------


def _build_context_block(recent_messages: list[str] | None) -> str:
    """Build conversation context block from recent messages (max 5, 300 chars each)."""
    if not recent_messages:
        return ""
    trimmed = [m[:300] for m in recent_messages[-5:]]
    return "\n\nRecent conversation context (for resolving pronouns and references):\n" + "\n".join(
        f"- {m}" for m in trimmed
    )


async def _entity_scan(
    llm: LLMProvider,
    message: str,
    recent_messages: list[str] | None = None,
    timeout_sec: float | None = None,
) -> list[ExtractedEntity]:
    """Pass 1: Discover all entities in the message."""
    context = _build_context_block(recent_messages)
    user_prompt = f"Current message to extract from:\n{message}{context}"

    raw = await _call_llm(
        llm=llm,
        system_prompt=ENTITY_SCAN_PROMPT,
        user_prompt=user_prompt,
        call_type="entity_scan",
        timeout_sec=timeout_sec,
    )

    if raw is None:
        return []

    try:
        cleaned = _strip_markdown_fences(raw)
        data = json.loads(cleaned)
        entities = [ExtractedEntity.model_validate(e) for e in (data.get("entities") or [])]
        logger.info("entity_scan: found %d entities", len(entities))
        return entities
    except Exception as e:
        logger.warning("entity_scan: parse failed: %s", e)
        return []


async def _extract_facts_for_batch(
    llm: LLMProvider,
    message: str,
    entity_names: list[str],
    recent_messages: list[str] | None = None,
    timeout_sec: float | None = None,
) -> list[ExtractedFact]:
    """Pass 2: Extract atomic facts for a batch of entities."""
    context = _build_context_block(recent_messages)
    entity_list = ", ".join(entity_names)
    user_prompt = (
        f"Current message to extract from:\n{message}{context}\n\n"
        f"Entities to extract facts for: {entity_list}\n"
        f"Extract ALL atomic facts about each of these entities from the message above."
    )

    raw = await _call_llm(
        llm=llm,
        system_prompt=FACT_EXTRACTION_PROMPT,
        user_prompt=user_prompt,
        call_type="fact_extraction",
        timeout_sec=timeout_sec,
    )

    if raw is None:
        return []

    try:
        cleaned = _strip_markdown_fences(raw)
        data = json.loads(cleaned)
        facts = [ExtractedFact.model_validate(f) for f in (data.get("facts") or [])]
        return facts
    except Exception as e:
        logger.warning("fact_extraction: parse failed for batch %s: %s", entity_names, e)
        return []


async def _extract_relations(
    llm: LLMProvider,
    message: str,
    entities: list[ExtractedEntity],
    timeout_sec: float | None = None,
) -> tuple[list[ExtractedRelation], bool]:
    """Pass 3: Extract relationships between entities.

    Retries once if the LLM returns 0 relations and there are 2+ entities
    (likely LLM variance, not a legitimate empty result).

    Returns (relations, retry_triggered).
    """
    entity_list = ", ".join(e.name for e in entities)
    user_prompt = (
        f"Current message to extract from:\n{message}\n\n"
        f"Known entities: {entity_list}\n"
        f"Extract ALL relationships between these entities from the message above."
    )

    relations = await _do_extract_relations(llm, user_prompt, timeout_sec)

    # Retry once if 0 relations with 2+ entities (LLM variance)
    retry_triggered = False
    if len(relations) == 0 and len(entities) >= 2:
        logger.info("relation_extraction: 0 relations with %d entities, retrying", len(entities))
        retry_triggered = True
        relations = await _do_extract_relations(llm, user_prompt, timeout_sec)

    logger.info("relation_extraction: found %d relations (retry=%s)", len(relations), retry_triggered)
    return relations, retry_triggered


async def _do_extract_relations(
    llm: LLMProvider,
    user_prompt: str,
    timeout_sec: float | None = None,
) -> list[ExtractedRelation]:
    """Single LLM call for relation extraction. Returns empty list on failure."""
    raw = await _call_llm(
        llm=llm,
        system_prompt=RELATION_EXTRACTION_PROMPT,
        user_prompt=user_prompt,
        call_type="relation_extraction",
        timeout_sec=timeout_sec,
    )

    if raw is None:
        return []

    try:
        cleaned = _strip_markdown_fences(raw)
        data = json.loads(cleaned)
        return [ExtractedRelation.model_validate(r) for r in (data.get("relations") or [])]
    except Exception as e:
        logger.warning("relation_extraction: parse failed: %s", e)
        return []


async def _multi_pass_extract(
    message: str,
    llm: LLMProvider,
    config: MemoryConfig,
    recent_messages: list[str] | None = None,
    timeout_sec: float | None = None,
) -> tuple[ExtractionOutput, dict]:
    """Multi-pass extraction: entity scan → facts per batch → relations.

    Falls back to single-pass if entity scan finds ≤ threshold entities.
    """
    # Pass 1: Entity scan
    entities = await _entity_scan(llm, message, recent_messages=recent_messages, timeout_sec=timeout_sec)

    # Pass 2: Extract facts (all entities in one call) + relations, concurrently
    entity_names = [e.name for e in entities]

    fact_task = _extract_facts_for_batch(
        llm, message, entity_names, recent_messages=recent_messages, timeout_sec=timeout_sec
    )
    relation_task = _extract_relations(llm, message, entities, timeout_sec=timeout_sec)

    fact_result, rel_result = await asyncio.gather(fact_task, relation_task)

    raw_facts: list[ExtractedFact] = list(fact_result)
    relations: list[ExtractedRelation] = rel_result[0]
    relation_retry_triggered: bool = rel_result[1]

    all_facts = _deduplicate_facts(raw_facts)

    logger.info(
        "extraction_multi_pass: entities=%d facts=%d (deduped from %d) relations=%d",
        len(entities),
        len(all_facts),
        len(raw_facts),
        len(relations),
    )

    extraction_meta = {
        "effective_mode": "multi_pass",
        "batch_count": 1,
        "relation_retry_triggered": relation_retry_triggered,
    }

    return _normalize_extraction_subjects(
        ExtractionOutput(
            entities=entities,
            facts=all_facts,
            relations=relations,
        )
    ), extraction_meta


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


async def extract_from_message(
    message: str,
    llm: LLMProvider,
    config: MemoryConfig | None = None,
    recent_messages: list[str] | None = None,
) -> tuple[ExtractionOutput, dict]:
    """Extract entities, facts, and relations from a user message.

    Args:
        message: The user message to extract from.
        llm: Injected LLM provider.
        config: Optional MemoryConfig for extraction tuning.
        recent_messages: Optional conversation context (last N messages)
            for resolving pronouns and anaphora references.

    Returns a tuple of (ExtractionOutput, extraction_meta dict).
    ExtractionOutput is empty on total failure (fail-safe).
    extraction_meta contains: effective_mode, batch_count, relation_retry_triggered.
    """
    if config is None:
        config = MemoryConfig()

    _timeout = config.extraction_timeout_sec if config.extraction_timeout_sec > 0 else None

    output, meta = await _multi_pass_extract(
        message, llm, config, recent_messages=recent_messages, timeout_sec=_timeout
    )

    logger.info(
        "extraction: entities=%d facts=%d relations=%d",
        len(output.entities),
        len(output.facts),
        len(output.relations),
    )

    return output, meta
