"""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, products, events, concepts, etc. mentioned.
- 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: 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 should be "place", not "person". A company should be "organization", not "place". A software product should be "product". An event should be "event". Cross-reference entity types with the facts.

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