"""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
import re

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 (as mentioned), type (person/organization/place/pet/other)
- The user themselves is always: name="user", type="person"

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.
- Confidence: 0.95 (explicitly stated), 0.80 (strongly implied), 0.60 (weakly implied). Skip below 0.5.

RELATIONS: Connections between entities.
- Each relation has: source, target, rel_type
- Allowed types: works_at, manages, reports_to, family_of, friend_of, partner_of, owns, lives_in, member_of, studies_at, works_with
- Example: "Kevin reports to Julia" → source="Kevin", target="Julia", rel_type="reports_to"

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

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 (as mentioned in the message), type (person/organization/place/pet/other)
- The user themselves is always included: name="user", type="person"
- Include EVERY named entity, even if mentioned only once.
- For nicknames/aliases, create ONE entity with the primary name.

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

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.

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
Allowed types: works_at, manages, reports_to, family_of, friend_of, partner_of, owns, lives_in, member_of, studies_at, works_with

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


def _strip_markdown_fences(raw: str) -> str:
    """Strip markdown code fences from LLM response."""
    if raw.startswith("```"):
        raw = re.sub(r"^```\w*\n?", "", raw)
        raw = re.sub(r"\n?```$", "", raw)
    return raw.strip()


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)


async def _call_llm(
    llm: LLMProvider,
    system_prompt: str,
    user_prompt: str,
    call_type: str = "extraction",
) -> str | None:
    """Call LLM with JSON mode. Returns raw content string or None on failure."""
    try:
        return await llm.complete(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0,
            response_format={"type": "json_object"},
        )
    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,
) -> 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",
    )

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

    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],
) -> list[ExtractedRelation]:
    """Pass 3: Extract relationships between entities."""
    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."
    )

    raw = await _call_llm(
        llm=llm,
        system_prompt=RELATION_EXTRACTION_PROMPT,
        user_prompt=user_prompt,
        call_type="relation_extraction",
    )

    if raw is None:
        return []

    try:
        cleaned = _strip_markdown_fences(raw)
        data = json.loads(cleaned)
        relations = [ExtractedRelation.model_validate(r) for r in (data.get("relations") or [])]
        logger.info("relation_extraction: found %d relations", len(relations))
        return relations
    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,
) -> ExtractionOutput:
    """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)

    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,
        )
        return await _single_pass_extract(message, llm, recent_messages=recent_messages)

    # 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) for batch in batches]
    relation_task = _extract_relations(llm, message, entities)

    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] = all_results[-1]  # type: ignore[assignment]

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

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


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


async def _single_pass_extract(
    message: str,
    llm: LLMProvider,
    recent_messages: list[str] | 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

    # Attempt 1
    try:
        raw = await llm.complete(
            messages=messages,
            temperature=0,
            response_format={"type": "json_object"},
        )
        output = _parse_extraction_output(raw)
    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:
            raw = await llm.complete(
                messages=messages,
                temperature=0,
                response_format={"type": "json_object"},
            )
            output = _parse_extraction_output(raw)
        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 output


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


async def _reflexion_pass(
    message: str,
    output: ExtractionOutput,
    llm: LLMProvider,
) -> 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",
    )

    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,
) -> ExtractionOutput:
    """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 ExtractionOutput (empty on total failure — fail-safe).
    """
    if config is None:
        config = MemoryConfig()

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

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

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

    return output
