src/agents/finalyzer_agent.py

"""
Finalyzer Agent: Creates final polished solution with clear output.

This agent runs when the verifier confirms the plan is sufficient.
It generates a final version of the code with improved formatting and output.
"""

from langchain_core.messages import AIMessage

from ..utils.code_execution import execute_code_safely
from ..utils.formatters import extract_code, format_data_descriptions, gemini_text
from ..utils.state import DSStarState


def finalyzer_node(state: DSStarState) -> dict:
    """
    Finalyzer Agent Node: Creates final polished solution.

    Takes the working code and creates a final version with:
    - Clear answer to the original question
    - Proper output formatting
    - Self-contained executable code

    Args:
        state: Current DSStarState

    Returns:
        Dictionary with updated state fields:
        - current_code: Final polished code
        - execution_result: Final execution output
        - messages: Agent communication messages
        - next: "__end__" (workflow complete)
    """
    print("=" * 60)
    print("FINALYZER AGENT STARTING...")
    print("=" * 60)

    data_context = format_data_descriptions(state["data_descriptions"])

    prompt = f"""You are an expert data analyst creating final solutions.

Original Question: {state["query"]}

Available Data:
{data_context}

Working Code:
{state["current_code"]}

Execution Result:
{state["execution_result"]}

Task: Create a final version of the code that:
1. Clearly prints the answer to the question
2. Includes proper formatting of the output
3. Is self-contained and executable
4. Has clear print statements

Provide ONLY the final Python code in a markdown code block."""

    try:
        # Get LLM response
        response = state["llm"].invoke(prompt)

        # Handle different response formats
        if hasattr(response, "content") and isinstance(response.content, list):
            response_text = gemini_text(response)
        elif hasattr(response, "content"):
            response_text = response.content
        else:
            response_text = str(response)

        final_code = extract_code(response_text)

        print("\nFinal Code Generated:")
        print("-" * 60)
        print(final_code[:300] + "..." if len(final_code) > 300 else final_code)
        print("-" * 60)

        # Execute final code
        print("\nExecuting final code...")
        success, stdout, stderr = execute_code_safely(final_code)

        if success:
            final_result = stdout
            print("\n✓ Final execution successful")
        else:
            # If final execution fails, use previous result
            print("\n⚠ Final execution failed, using previous result")
            final_result = state["execution_result"]

        print("\nFinal Result:")
        print("-" * 60)
        print(final_result[:300] + "..." if len(final_result) > 300 else final_result)
        print("-" * 60)
        print("=" * 60)
        print("SOLUTION COMPLETE ✓")
        print("=" * 60)

        return {
            "current_code": final_code,
            "execution_result": final_result,
            "messages": [AIMessage(content="Solution finalized")],
            "next": "__end__",
        }

    except Exception as e:
        # On error, return current state
        print(f"\n✗ Finalyzer error: {str(e)}")
        print("Using current solution as final")
        return {
            "messages": [
                AIMessage(content=f"Finalyzer error: {str(e)}, using current solution")
            ],
            "next": "__end__",
        }


# Standalone test function
def test_finalyzer(
    llm, query: str, data_descriptions: dict, current_code: str, execution_result: str
):
    """
    Test the finalyzer agent independently.

    Args:
        llm: LLM instance
        query: User query
        data_descriptions: Dict of filename -> description
        current_code: Working code to finalize
        execution_result: Current execution result

    Returns:
        Dictionary with finalyzer results
    """
    # Create minimal test state
    test_state = {
        "llm": llm,
        "query": query,
        "data_descriptions": data_descriptions,
        "plan": [],
        "current_code": current_code,
        "execution_result": execution_result,
        "is_sufficient": True,
        "router_decision": "",
        "iteration": 0,
        "max_iterations": 20,
        "messages": [],
        "next": "finalyzer",
    }

    result = finalyzer_node(test_state)

    print("\n" + "=" * 60)
    print("FINALYZER TEST RESULTS")
    print("=" * 60)
    print("Final Code:")
    print(result.get("current_code", "No code"))
    print("\nFinal Result:")
    print(result.get("execution_result", "No result"))

    return result