# auditor.py
from langchain.agents import create_openai_tools_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool, ToolException
from langchain_core.messages import AIMessage
import json
import os
import re
from state import AgentState
from db_utils import extract_and_clean_json
from terminal_utils import print_colorful_break

async def auditor_agent_node(state: AgentState):
    print_colorful_break("AUDITOR")
    logger = state['logger']
    logger.info("--- Running Auditor Agent ---")
    all_tools = state['mcp_tools']
    model = state['model']
    timestamp = state['timestamp']
    
    auditor_tools = [t for t in all_tools if t.name in ["graphs_get", "query"]]
    auditor_prompt = '''Your goal is to review the LightRAG knowledge base for data quality issues.
1.  **Identify issues**: Scan the graph for duplicates, irregular normalization, etc.
2.  **Generate Report**: Create a report of your findings.
3.  **Save Report**: Use `save_report` to save the findings to `auditor_report.json`.'''

    prompt = ChatPromptTemplate.from_template(auditor_prompt)
    agent_executor = create_openai_tools_agent(model, auditor_tools, prompt)

    task_input = "Your task is to audit the knowledge base. Begin now."

    result = await agent_executor.ainvoke({"input": task_input, "timestamp": timestamp})
    
    logger.info(f"Auditor Agent finished with output: {result['output']}")
    return {"messages": state['messages'] + [AIMessage(content=result['output'])]}

def save_auditor_report_node(state: AgentState):
    """Saves the final report from the last AI message."""
    logger = state['logger']
    final_message_from_agent = state['messages'][-1]

    status = f"--- Saving Auditor Report ---\n{final_message_from_agent.content}"
    logger.info(status)

    try:
        report_json = extract_and_clean_json(final_message_from_agent.content)
        if 'report_id' not in report_json:
            report_json['report_id'] = state.get('auditor_report_id', 'unknown_id')
        save_auditor_report({"auditor_report": json.dumps(report_json)})
        status = f"Successfully saved auditor report with ID {report_json.get('report_id')}"
        logger.info(status)
    except (ValueError, KeyError) as e:
        status = f"Error processing or saving auditor report: {e}"
        logger.error(status, exc_info=True)

    return {"messages": state['messages'] + [AIMessage(content=status)]}