import { StateGraph, Annotation } from '@langchain/langgraph';
import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
import { HumanMessage, SystemMessage } from '@langchain/core/messages';
import type { FastifyBaseLogger } from 'fastify';

/**
 * State for strategy analysis workflow
 */
const StrategyAnalysisState = Annotation.Root({
  strategyCode: Annotation<string>(),
  ticker: Annotation<string>(),
  timeframe: Annotation<string>(),

  // Analysis steps
  codeReview: Annotation<string | null>({
    default: () => null,
  }),
  backtestResults: Annotation<Record<string, unknown> | null>({
    default: () => null,
  }),
  riskAssessment: Annotation<string | null>({
    default: () => null,
  }),
  humanApproved: Annotation<boolean>({
    default: () => false,
  }),

  // Final output
  recommendation: Annotation<string | null>({
    default: () => null,
  }),
});

type StrategyAnalysisStateType = typeof StrategyAnalysisState.State;

/**
 * Build strategy analysis workflow using LangGraph
 *
 * Workflow steps:
 * 1. Code review (LLM analyzes strategy code)
 * 2. Backtest (calls user's MCP backtest tool)
 * 3. Risk assessment (LLM evaluates results)
 * 4. Human approval (pause for user review)
 * 5. Final recommendation
 */
export function buildStrategyAnalysisWorkflow(
  model: BaseChatModel,
  logger: FastifyBaseLogger,
  mcpBacktestFn: (strategy: string, ticker: string, timeframe: string) => Promise<Record<string, unknown>>
) {
  // Node: Code Review
  const codeReviewNode = async (state: StrategyAnalysisStateType) => {
    logger.info('Strategy workflow: Code review');

    const systemPrompt = `You are an expert trading strategy analyst.
Review the following strategy code for potential issues, bugs, or improvements.
Focus on: logic errors, edge cases, performance, and trading best practices.`;

    const response = await model.invoke([
      new SystemMessage(systemPrompt),
      new HumanMessage(`Review this strategy:\n\n${state.strategyCode}`),
    ]);

    return {
      codeReview: response.content as string,
    };
  };

  // Node: Backtest
  const backtestNode = async (state: StrategyAnalysisStateType) => {
    logger.info('Strategy workflow: Running backtest');

    const results = await mcpBacktestFn(state.strategyCode, state.ticker, state.timeframe);

    return {
      backtestResults: results,
    };
  };

  // Node: Risk Assessment
  const riskAssessmentNode = async (state: StrategyAnalysisStateType) => {
    logger.info('Strategy workflow: Risk assessment');

    const systemPrompt = `You are a risk management expert for trading strategies.
Analyze the backtest results and provide a risk assessment.
Focus on: drawdown, win rate, Sharpe ratio, position sizing, and risk of ruin.`;

    const response = await model.invoke([
      new SystemMessage(systemPrompt),
      new HumanMessage(
        `Code review: ${state.codeReview}\n\nBacktest results: ${JSON.stringify(state.backtestResults, null, 2)}\n\nProvide risk assessment:`
      ),
    ]);

    return {
      riskAssessment: response.content as string,
    };
  };

  // Node: Human Approval (placeholder - would integrate with UI)
  const humanApprovalNode = async (state: StrategyAnalysisStateType) => {
    logger.info('Strategy workflow: Awaiting human approval');

    // In real implementation, this would pause and wait for user input
    // For now, auto-approve
    return {
      humanApproved: true,
    };
  };

  // Node: Final Recommendation
  const recommendationNode = async (state: StrategyAnalysisStateType) => {
    logger.info('Strategy workflow: Generating recommendation');

    const systemPrompt = `Provide a final recommendation on whether to deploy this trading strategy.
Summarize the code review, backtest results, and risk assessment.
Give clear go/no-go decision with reasoning.`;

    const response = await model.invoke([
      new SystemMessage(systemPrompt),
      new HumanMessage(
        `Code review: ${state.codeReview}\n\nBacktest: ${JSON.stringify(state.backtestResults)}\n\nRisk: ${state.riskAssessment}\n\nApproved: ${state.humanApproved}\n\nYour recommendation:`
      ),
    ]);

    return {
      recommendation: response.content as string,
    };
  };

  // Build graph
  const workflow = new StateGraph(StrategyAnalysisState)
    .addNode('code_review', codeReviewNode)
    .addNode('backtest', backtestNode)
    .addNode('risk_assessment', riskAssessmentNode)
    .addNode('human_approval', humanApprovalNode)
    .addNode('recommendation', recommendationNode)
    .addEdge('__start__', 'code_review')
    .addEdge('code_review', 'backtest')
    .addEdge('backtest', 'risk_assessment')
    .addEdge('risk_assessment', 'human_approval')
    .addConditionalEdges('human_approval', (state) => {
      return state.humanApproved ? 'recommendation' : '__end__';
    })
    .addEdge('recommendation', '__end__');

  return workflow.compile();
}

/**
 * Example usage:
 *
 * const workflow = buildStrategyAnalysisWorkflow(model, logger, mcpBacktestFn);
 *
 * const result = await workflow.invoke({
 *   strategyCode: "strategy code here",
 *   ticker: "BTC/USDT",
 *   timeframe: "1h",
 * });
 *
 * console.log(result.recommendation);
 */
