"""
{{ display_name }} - Extract node.

This module handles entity and insight extraction using LLM.
"""

import json

from torivers_sdk.automation.output import JsonDataBlock, build_output_data
from torivers_sdk.tools import LLMClient, LLMMessage

from state import {{ class_name }}State

{% raw %}
EXTRACT_PROMPT = """Extract key entities and insights from the following content.

Content:
{content}

Summary (for context):
{summary}

Return a JSON object with:
{{
  "entities": [
    {{"name": "...", "type": "person|organization|location|product|other", "relevance": 0.0-1.0}}
  ],
  "insights": [
    {{"insight": "...", "importance": "high|medium|low", "category": "..."}}
  ]
}}

Return only valid JSON, no additional text."""
{% endraw %}


async def extract(state: {{ class_name }}State) -> dict:
    """Extract entities and insights from content.

    Args:
        state: Current analysis state.

    Returns:
        Dictionary with state updates including entities and insights.
    """
    context = state["context"]
    input_data = state["input_data"]

    content = input_data.get("content", "")
    summary = state.get("summary", "")

    context.log_progress("Extracting", "Extracting entities and insights...")

    llm = LLMClient.get_default()
    response = await llm.chat(
        messages=[
            LLMMessage(
                role="system",
                content="You are an expert at extracting structured information. Return only valid JSON.",
            ),
            LLMMessage(
                role="user",
                content=EXTRACT_PROMPT.format(content=content, summary=summary),
            ),
        ],
    )

    # Parse JSON response
    try:
        extraction = json.loads(response.content)
        entities = extraction.get("entities", [])
        insights = extraction.get("insights", [])
    except json.JSONDecodeError:
        entities = []
        insights = []

    total_tokens = state.get("total_tokens_used", 0) + response.tokens_used
    context.log_success(
        "Extraction complete",
        f"Found {len(entities)} entities, {len(insights)} insights",
    )

    return {
        "current_step": "completed",
        "entities": entities,
        "insights": insights,
        "total_tokens_used": total_tokens,
        "output_data": build_output_data(
            JsonDataBlock(
                label="Analysis Results",
                data={
                    "summary": state.get("summary"),
                    "sentiment": state.get("sentiment"),
                    "entities": entities,
                    "insights": insights,
                },
                display_hint="key_value",
            ),
        ),
        "metadata": {
            "extraction_tokens": response.tokens_used,
            "total_tokens": total_tokens,
        },
    }
