from seer import SeerClient
from opentelemetry import trace

client = SeerClient()
tracer = trace.get_tracer(__name__)

def answer_multi_hop_question(query: str) -> str:
    """
    Example: "What awards did the director of Inception win?"

    This requires decomposing into:
    1. "Who directed Inception?"
    2. "What awards has [director] won?"
    """
    with tracer.start_as_current_span("multi_hop_query"):

        # Hop 1: Find the director
        # A query rewriter rewrites the original question to get the first piece
        subquery1 = "Who directed Inception?"

        with tracer.start_as_current_span("retrieval_hop_1"):
            hop1_context = retrieve(subquery1)

            client.log(
                task=query,                      # Original: "What awards did..."
                context=hop1_context,
                subquery=subquery1,              # "Who directed Inception?"
                span_name="retrieval_hop_1",
            )

            # Extract the answer: "Christopher Nolan"
            director = extract_entity(hop1_context, "director")

        # Hop 2: Find awards for the director
        # Query rewriter uses extracted entity to form the next subquery
        subquery2 = f"What awards has {director} won?"

        with tracer.start_as_current_span("retrieval_hop_2"):
            hop2_context = retrieve(subquery2)

            client.log(
                task=query,                      # Still the original query
                context=hop2_context,
                subquery=subquery2,              # "What awards has Christopher Nolan won?"
                span_name="retrieval_hop_2",
            )

        # Combine contexts from all hops
        all_context = hop1_context + hop2_context

        # Log the final joined context that goes to the LLM
        with tracer.start_as_current_span("context_join"):
            client.log(
                task=query,
                context=all_context,             # Combined context from all hops
                span_name="final_context",
                is_final_context=True,           # THIS is what the LLM sees
            )

        return synthesize_answer(query, all_context)

def agent_loop(query: str, max_iterations: int = 5):
    """Agent decides when to retrieve and when to answer."""
    context_so_far = []

    with tracer.start_as_current_span("agent_query"):
        for i in range(max_iterations):
            # Agent decides next action
            action = agent.plan_next_action(query, context_so_far)

            if action.type == "retrieve":
                # Agent wants more information
                results = retrieve(action.search_query)
                context_so_far.extend(results)

                client.log(
                    task=query,
                    context=results,
                    subquery=action.search_query,  # What the agent searched for
                    span_name=f"agent_retrieval_{i}",
                    metadata={
                        "iteration": i,
                        "agent_reasoning": action.reasoning,
                    },
                )

            elif action.type == "answer":
                # Agent is ready to synthesize
                # Mark the final context
                client.log(
                    task=query,
                    context=context_so_far,
                    span_name="final_context",
                    is_final_context=True,
                )
                break

        return agent.synthesize(query, context_so_far)

with tracer.start_as_current_span("parallel_retrieval"):
    # Search multiple sources simultaneously
    wiki_results = retrieve_from_wiki(query)
    kb_results = retrieve_from_kb(query)

    client.log(task=query, context=wiki_results, span_name="retrieval_wiki")
    client.log(task=query, context=kb_results, span_name="retrieval_kb")

    # Combine and pass to LLM
    combined = wiki_results + kb_results
    client.log(
        task=query,
        context=combined,
        span_name="final_merged",
        is_final_context=True,
    )

# Initial retrieval
initial_context = retrieve(query)
client.log(task=query, context=initial_context, span_name="retrieval_initial")

# LLM suggests refinement
refined_query = llm.suggest_refined_query(query, initial_context)

# Refined retrieval
refined_context = retrieve(refined_query)
client.log(
    task=query,
    context=refined_context,
    subquery=refined_query,
    span_name="retrieval_refined",
    is_final_context=True,
)

# The context that actually goes to the LLM
client.log(..., is_final_context=True)

# Correct: Same task, different subqueries
client.log(task=original_query, subquery="Who is X?", ...)
client.log(task=original_query, subquery="What did X do?", ...)

# Wrong: Changing task per hop
client.log(task="Who is X?", ...)  # Don't do this

# If task recall is low but subquery recall is high,
# the problem is query decomposition, not retrieval