import { ChatOpenAI } from "@langchain/openai";
import { SystemMessage, HumanMessage, AIMessage } from "@langchain/core/messages";
import { RunnableConfig } from "@langchain/core/runnables";
import { EvaluationState } from "../state";

/**
 * Node responsible for grading the user's answer and deciding if a follow-up is needed.
 */
export const graderNode = async (
    state: EvaluationState,
    config?: RunnableConfig
): Promise<Partial<EvaluationState>> => {
    const { model } = (config?.configurable as any) || {};
    const { questions, currentQuestionIndex, messages } = state;

    console.log("[GraderNode] Entering node...", {
        currentIndex: currentQuestionIndex,
        numMessages: messages?.length
    });

    if (!model) {
        throw new Error("Missing model in node configuration");
    }

    const currentQuestion = questions[currentQuestionIndex];
    const lastUserMessage = messages[messages.length - 1];

    if (!(lastUserMessage instanceof HumanMessage)) {
        return {};
    }

    const systemPrompt = `You are an expert examiner. 
Grade the user's answer based on the following question and key points.
IMPORTANT: You MUST provide the feedback in the following language: ${state.language || 'zh'}.

QUESTION: ${currentQuestion.questionText}
EXPECTED KEY POINTS: ${currentQuestion.keyPoints.join(", ")}

Evaluate:
1. Accuracy: Did they cover the key points correctly?
2. Completeness: Did they miss anything important?
3. Depth: Is the explanation sufficient?

Provide:
1. A score from 0 to 10.
2. Constructive feedback.
3. A boolean flag 'should_follow_up' if the answer is incomplete or unclear and needs further clarification.

Format your response as JSON:
{
  "score": 8,
  "feedback": "...",
  "should_follow_up": false
}`;

    const response = await model.invoke([
        new SystemMessage(systemPrompt),
        new HumanMessage((lastUserMessage as HumanMessage).content as string),
    ]);

    try {
        const result = JSON.parse(response.content as string);
        console.log("[GraderNode] AI Grade Result:", result);
        const isZh = state.language === 'zh';
        const isJa = state.language === 'ja';
        const scoreLabel = isZh ? '得分' : isJa ? 'スコア' : 'Score';
        const feedbackLabel = isZh ? '反馈' : isJa ? 'フィードバック' : 'Feedback';

        const feedbackMessage = new AIMessage(`${scoreLabel}: ${result.score}/10\n\n${feedbackLabel}: ${result.feedback}`);

        const newScores = { ...state.scores, [currentQuestion.id || currentQuestionIndex]: result.score };
        
        let shouldFollowUp = result.should_follow_up;
        const currentFollowUpCount = state.followUpCount || 0;

        // Breakout logic: 
        // 1. Max 1 follow-up per question
        // 2. If score is decent (>= 8), don't follow up
        // 3. If answer is short "don't know", don't follow up
        const userContent = (lastUserMessage.content as string).trim().toLowerCase();
        const saysIDontKnow = userContent.length < 10 && (
            userContent.includes("不知道") || 
            userContent.includes("不会") || 
            userContent.includes("don't know") || 
            userContent.includes("no idea")
        );

        if (currentFollowUpCount >= 1 || result.score >= 8 || saysIDontKnow) {
            shouldFollowUp = false;
        }

        return {
            feedbackHistory: [feedbackMessage],
            scores: newScores,
            shouldFollowUp: shouldFollowUp,
            followUpCount: shouldFollowUp ? currentFollowUpCount + 1 : 0,
            currentQuestionIndex: shouldFollowUp ? currentQuestionIndex : currentQuestionIndex + 1,
        };
    } catch (error) {
        console.error("Failed to parse grade from AI response:", error);
        return {
            feedbackHistory: [new AIMessage("I had some trouble grading that, but let's move on.")],
            currentQuestionIndex: currentQuestionIndex + 1,
        };
    }
};