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

Supports two modes (configurable via MemoryConfig.extraction_mode):
- "multi_pass": Entity scan → per-entity fact extraction → relation extraction.
  Better for dense messages with many entities. Default.
- "single_pass": One LLM call for everything. Faster/cheaper, good for simple messages.

The multi-pass pipeline automatically falls back to single-pass when the
entity count is low (≤ config.multi_pass_entity_threshold).

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
# ---------------------------------------------------------------------------
EXTRACTION_SYSTEM_PROMPT = """You are a personal knowledge extractor.
Given a conversation message, extract:

ENTITIES: People, places, organizations, pets, etc. mentioned.
- Each entity has: name, type (person/organization/place/pet/other), aliases (optional list of alternative names)
- The user themselves is always: name="user", type="person"
- IMPORTANT: When the same entity is referred to by multiple names (nickname + full name, abbreviation + full form, colloquial + formal), 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"]}
  - "I live in Sampa, São Paulo is great" → {"name": "São Paulo", "type": "place", "aliases": ["Sampa"]}
  - "my cat Bokita (sometimes I write Boquita)" → {"name": "Bokita", "type": "pet", "aliases": ["Boquita"]}
- Only group names when the message explicitly indicates they refer to the same entity. Do not guess.

FACTS: Factual information about entities.
- Each fact is a short natural language sentence.
- Include: subject (who/what the fact is about), fact (the information), confidence (0.0-1.0)
- Use "user" as subject when the message is about the user themselves.
- Use the entity's primary name (not aliases) as the subject.
- Confidence: 0.95 (explicitly stated), 0.80 (strongly implied), 0.60 (weakly implied). Skip below 0.5.
- When a relationship is strongly implied but not directly stated as a fact, also extract it as a fact. Example: "I'm going to Curitiba to visit my mom" implies → fact: "User's mom lives in Curitiba" (confidence: 0.60).
- This ensures implied knowledge has both a graph edge AND a searchable fact.

RELATIONS: Connections between entities.
- Each relation has: source, target, rel_type
- Use the entity's primary name (not aliases) as source/target.
- Use a short, descriptive snake_case type (e.g., works_at, manages, reports_to, lives_in, mentored_by, inspired_by, competed_with).
- Common types: works_at, manages, reports_to, family_of, friend_of, partner_of, owns, lives_in, member_of, studies_at, works_with.
- You may use other types when they better describe the relationship.
- Example: "Kevin reports to Julia" → source="Kevin", target="Julia", rel_type="reports_to"
- Do NOT create same_as/also_known_as relations for entities you already grouped with aliases.

Rules:
- Only extract what was said or strongly implied. Never invent.
- Prefer atomic facts (1 piece of information per fact).
- If nothing to extract, return empty lists.
- Extract ALL named people/places/organizations as entities, even if no facts about them.

Respond in JSON with keys: entities, facts, relations."""

REFLEXION_SYSTEM_PROMPT = """You are a quality reviewer for knowledge extraction.
Review the extraction output for a conversation message.

Check:
1. Are there people mentioned in the message who are missing from entities?
2. Should any fact actually be a relation (e.g., "X works at Y" → relation, not fact)?
3. Are there duplicate or near-duplicate facts?
4. Are all entity names correct (not truncated or misspelled)?
5. Are there separate entities that clearly refer to the same real-world entity (e.g., a nickname and a full name)? If so, merge them into ONE entity using the most complete name as "name" and the other names in "aliases".
6. Is the type of each entity correct? A city, country, or neighborhood cannot be "person". A company cannot be "place". Cross-reference entity types with the facts — if a fact says "X is a neighborhood in Y", then Y must be type "place", not "person". Common mistake: when many people are mentioned alongside places, the LLM may misclassify places as "person".

If corrections are needed, return the corrected output in the same JSON format.
If no corrections are needed, return the same output unchanged."""

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

ENTITIES: Every person, place, organization, pet, product, team, squad, group, or named thing.
- Each entity has: name, type (person/organization/place/pet/other), 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 about each entity.

Rules:
- ONE fact = ONE piece of information. Never combine multiple pieces into one fact.
- Extract: roles, titles, responsibilities, attributes, preferences, locations, relationships, nicknames, scopes, affiliations.
- Include: subject (entity name), fact (natural language sentence), confidence (0.0-1.0)
- Use "user" as subject for facts about the user themselves.
- Confidence: 0.95 (explicitly stated), 0.80 (strongly implied), 0.60 (weakly implied).
- Do NOT skip any information. If someone has a role AND a scope AND a nickname, that is 3 separate facts.
- Organizational structure (which squad/team belongs to which tribe/division) = separate facts.
- Reporting relationships (who reports to whom, direct vs indirect) = separate facts.
- When a relationship is strongly implied but not directly stated as a fact, also extract it as a fact. Example: "I'm going to Curitiba to visit my mom" implies → fact: "User's mom lives in Curitiba" (confidence: 0.60). This ensures implied knowledge has both a graph edge AND a searchable fact.

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, manages, reports_to, lives_in, mentored_by, inspired_by, competed_with).
- Common types: works_at, manages, reports_to, family_of, friend_of, partner_of, owns, lives_in, member_of, studies_at, works_with.
- You may use other types when they better describe the relationship.

Rules:
- Extract EVERY relationship mentioned or strongly implied.
- Reporting chains: "A reports to B" → source="A", target="B", rel_type="reports_to"
- Management: "A manages B" → source="A", target="B", rel_type="manages"
- Membership: "squad X belongs to tribe Y" → source="squad X", target="tribe Y", rel_type="member_of"
- If A is a direct report and B is indirect (via C), extract both: A→user reports_to AND A→C reports_to.

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


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)

    if len(entities) <= config.multi_pass_entity_threshold:
        logger.info(
            "extraction: %d entities (≤ %d), falling back to single-pass",
            len(entities),
            config.multi_pass_entity_threshold,
        )
        output = await _single_pass_extract(message, llm, recent_messages=recent_messages, timeout_sec=timeout_sec)
        return output, {"effective_mode": "single_pass", "batch_count": 0, "relation_retry_triggered": False}

    # Pass 2: Extract facts in batches
    entity_names = [e.name for e in entities]
    batches = [
        entity_names[i : i + config.entity_batch_size] for i in range(0, len(entity_names), config.entity_batch_size)
    ]

    # Run fact extraction batches concurrently with relation extraction
    fact_tasks = [
        _extract_facts_for_batch(llm, message, batch, recent_messages=recent_messages, timeout_sec=timeout_sec)
        for batch in batches
    ]
    relation_task = _extract_relations(llm, message, entities, timeout_sec=timeout_sec)

    all_results = await asyncio.gather(*fact_tasks, relation_task)

    # Merge facts from all batches
    all_facts: list[ExtractedFact] = []
    for result in all_results[:-1]:  # all but last (relations)
        all_facts.extend(result)  # type: ignore[arg-type]

    relations: list[ExtractedRelation]
    relation_retry_triggered: bool
    _rel_result: tuple[list[ExtractedRelation], bool] = all_results[-1]  # type: ignore[assignment]
    relations = _rel_result[0]
    relation_retry_triggered = _rel_result[1]

    # Deduplicate facts (same subject + same fact text)
    seen = set()
    unique_facts = []
    for f in all_facts:
        key = (f.subject.lower().strip(), f.fact.lower().strip())
        if key not in seen:
            seen.add(key)
            unique_facts.append(f)

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

    extraction_meta = {
        "effective_mode": "multi_pass",
        "batch_count": len(batches),
        "relation_retry_triggered": relation_retry_triggered,
    }

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


# ---------------------------------------------------------------------------
# Single-pass extraction
# ---------------------------------------------------------------------------


async def _single_pass_extract(
    message: str,
    llm: LLMProvider,
    recent_messages: list[str] | None = None,
    timeout_sec: float | None = None,
) -> ExtractionOutput:
    """Single-pass extraction: one LLM call for entities + facts + relations."""
    context = _build_context_block(recent_messages)
    messages = [
        {"role": "system", "content": EXTRACTION_SYSTEM_PROMPT},
        {"role": "user", "content": f"Current message to extract from:\n{message}{context}"},
    ]

    output: ExtractionOutput | None = None

    async def _do_call() -> str:
        result = await llm.complete(
            messages=messages,
            temperature=0,
            response_format={"type": "json_object"},
        )
        return result.text

    # Attempt 1
    try:
        coro = _do_call()
        raw = await (asyncio.wait_for(coro, timeout=timeout_sec) if timeout_sec and timeout_sec > 0 else coro)
        output = _parse_extraction_output(raw)
    except TimeoutError:
        logger.warning("extraction: LLM call timed out (attempt 1)")
    except Exception as e:
        logger.warning("extraction: parse failed (attempt 1): %s", e)

    # Retry 1x on failure
    if output is None:
        logger.info("extraction: retrying")
        try:
            coro = _do_call()
            raw = await (asyncio.wait_for(coro, timeout=timeout_sec) if timeout_sec and timeout_sec > 0 else coro)
            output = _parse_extraction_output(raw)
        except TimeoutError:
            logger.warning("extraction: LLM call timed out (attempt 2)")
        except Exception as e:
            logger.warning("extraction: parse failed (attempt 2): %s", e)

    if output is None:
        logger.warning("extraction: total failure, returning empty")
        return ExtractionOutput()

    return _normalize_extraction_subjects(output)


# ---------------------------------------------------------------------------
# Reflexion pass (shared by both modes)
# ---------------------------------------------------------------------------


async def _reflexion_pass(
    message: str,
    output: ExtractionOutput,
    llm: LLMProvider,
    timeout_sec: float | None = None,
) -> ExtractionOutput:
    """Quality review pass — validates completeness, not domain rules.

    Fails silently (returns original output).
    """
    user_prompt = (
        f"Original message:\n{message}\n\n"
        f"Extraction output:\n{output.model_dump_json()}\n\n"
        "Review and correct if needed. Return JSON with entities, facts, relations."
    )

    raw = await _call_llm(
        llm=llm,
        system_prompt=REFLEXION_SYSTEM_PROMPT,
        user_prompt=user_prompt,
        call_type="reflexion",
        timeout_sec=timeout_sec,
    )

    if raw is None:
        return output

    try:
        corrected = _parse_extraction_output(raw)
        # Only accept if reflexion didn't empty everything
        if corrected.entities or corrected.facts or corrected.relations:
            return corrected
        return output
    except Exception:
        logger.debug("extraction: reflexion parse failed, keeping original")
        return output


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


async def extract_from_message(
    message: str,
    llm: LLMProvider,
    model: str = "gpt-4o",
    extraction_mode: str = "multi_pass",
    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.
        model: Model name for extraction (default: gpt-4o).
        extraction_mode: "multi_pass" (default) or "single_pass".
        config: Optional MemoryConfig for multi-pass 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

    default_meta = {"effective_mode": extraction_mode, "batch_count": 0, "relation_retry_triggered": False}

    if extraction_mode == "multi_pass":
        output, meta = await _multi_pass_extract(
            message, llm, config, recent_messages=recent_messages, timeout_sec=_timeout
        )
    else:
        output = await _single_pass_extract(message, llm, recent_messages=recent_messages, timeout_sec=_timeout)
        meta = default_meta

    # Reflexion pass (only for single-pass; multi-pass already has per-entity focus)
    if extraction_mode != "multi_pass" and (output.facts or output.entities):
        _ref_timeout = config.reflexion_timeout_sec if config.reflexion_timeout_sec > 0 else None
        output = await _reflexion_pass(message, output, llm, timeout_sec=_ref_timeout)

    logger.info(
        "extraction: mode=%s effective_mode=%s entities=%d facts=%d relations=%d",
        extraction_mode,
        meta.get("effective_mode", extraction_mode),
        len(output.entities),
        len(output.facts),
        len(output.relations),
    )

    return output, meta
