AITrader/app/lib/execution.ts

import type { TradingDecision, ExecutionPlan } from "../types/agents";

export function enrichExecutionPlan(decision: TradingDecision, input: any): TradingDecision {
  try {
    const prices: number[] = input?.technicalData?.prices || [];
    const entryPrice = prices.length ? prices[prices.length - 1] : undefined;

    // ATR approximation: prefer bar-based ATR (high-low average), fall back to price diffs
    let atr = 0;
    const bars: any[] = input?.technicalData?.bars || [];
    if (bars && bars.length >= 2) {
      let sum = 0;
      for (const b of bars) {
        const high = typeof b.HighPrice === 'number' ? b.HighPrice : (typeof b.h === 'number' ? b.h : 0);
        const low = typeof b.LowPrice === 'number' ? b.LowPrice : (typeof b.l === 'number' ? b.l : 0);
        sum += Math.max(0, high - low);
      }
      atr = sum / bars.length;
    } else if (prices && prices.length >= 2) {
      let sum = 0;
      for (let i = 1; i < prices.length; i++) sum += Math.abs(prices[i] - prices[i - 1]);
      atr = sum / (prices.length - 1);
    } else if (entryPrice) {
      atr = entryPrice * 0.01; // fallback 1%
    }

    const rr = 2; // default risk:reward
    const equity = Number(input?.account?.cash ?? input?.account?.buying_power ?? process.env.DEFAULT_ACCOUNT_EQUITY ?? 10000);

    if (!decision.executionPlan) decision.executionPlan = {} as ExecutionPlan;
    const plan = decision.executionPlan as any;

    const maxLossPercent = plan.riskManagement?.maxLossPercent ?? plan.maxLossPercent ?? 1; // default 1%
    const riskPercent = Number(maxLossPercent) / 100;

    // compute stop distance (price units)
    let stopDistanceByPercent = entryPrice ? Math.abs(entryPrice * riskPercent) : 0;
    const stopDistanceByAtr = atr ? atr * 1.5 : 0; // multiplier
    let stopDistance = Math.max(stopDistanceByAtr, stopDistanceByPercent, 0.0001);

    // compute stopLoss absolute price if missing
    if (plan.stopLoss == null && entryPrice != null) {
      if (decision.action === 'buy') {
        plan.stopLoss = Number((entryPrice - stopDistance).toFixed(2));
      } else if (decision.action === 'sell') {
        // for sell (exit/short) place stop above entry
        plan.stopLoss = Number((entryPrice + stopDistance).toFixed(2));
      } else {
        // default: buy-style stop
        plan.stopLoss = Number((entryPrice - stopDistance).toFixed(2));
      }
    }

    // compute takeProfit if missing
    if (plan.takeProfit == null && entryPrice != null) {
      if (decision.action === 'buy') {
        plan.takeProfit = Number((entryPrice + stopDistance * rr).toFixed(2));
      } else if (decision.action === 'sell') {
        plan.takeProfit = Number((entryPrice - stopDistance * rr).toFixed(2));
      } else {
        plan.takeProfit = Number((entryPrice + stopDistance * rr).toFixed(2));
      }
    }

    // compute shares if missing using risk-based sizing
    if (plan.amount == null) {
      const riskAmount = equity * riskPercent;
      // Protect against extremely small stopDistance (which can occur with missing/flat prices)
      if (!stopDistance || stopDistance < 0.01) {
        stopDistance = entryPrice ? Math.max(entryPrice * 0.01, 0.01) : 1; // default 1 unit when no price
      }
      const rawShares = Math.max(1, Math.floor(riskAmount / stopDistance));
      // Cap shares to what is affordable with full equity and a reasonable absolute cap
      const affordableMax = Math.max(1, Math.floor(equity / Math.max(entryPrice || 1, 1)));
      const absoluteMax = 100000; // safety cap
      const shares = Math.min(rawShares, affordableMax, absoluteMax);
      plan.amount = shares;
    }

    // normalize nested riskManagement
    plan.riskManagement = plan.riskManagement || {};
    if (plan.riskManagement.maxLossPercent == null) plan.riskManagement.maxLossPercent = maxLossPercent;

    decision.executionPlan = plan as ExecutionPlan;
  } catch (err) {
    console.warn("enrichExecutionPlan error:", err);
  }
  return decision;
}

// Optional LLM verification step: review computed executionPlan and suggest adjustments
export async function verifyExecutionPlanWithLLM(decision: TradingDecision, input: any, model?: string): Promise<TradingDecision> {
  try {
    const apiKey = process.env.OPENROUTER_API_KEY;
    if (!apiKey) return decision;

    const { OpenRouterClient } = await import("./openrouter");
    const client = new OpenRouterClient(apiKey, { defaultModel: model });

    const plan = decision.executionPlan || {};
    const prices: number[] = input?.technicalData?.prices || [];
    const entryPrice = prices.length ? prices[prices.length - 1] : null;

    const userPrompt = `Review the following deterministic execution plan and either approve or provide corrected values. Respond with a JSON object only with the shape: { "approved": boolean, "executionPlan": { /* fields to use */ }, "notes": "..." }\n\nContext:\nentryPrice: ${entryPrice}\nrecentPrices: ${JSON.stringify(prices.slice(-10))}\nComputedPlan: ${JSON.stringify(plan)}\n\nGuidelines:\n- If values look reasonable, return {"approved": true, "executionPlan": {}} (empty executionPlan means keep computed values).\n- If any value should be adjusted, return an executionPlan object with corrected fields (amount, stopLoss, takeProfit, riskManagement).\n- Do not include any extra text outside the JSON object.`;

    const messages = [
      { role: "system", content: "You are a conservative trading assistant. Validate risk sizing and stop levels. If you suggest changes, prefer conservative (smaller) position sizes and wider stops only if volatility justifies it." },
      { role: "user", content: userPrompt },
    ];

    const res: any = await client.createChatCompletion(messages as any);
    const content = res?.choices?.[0]?.message?.content ?? "";

    // try to parse JSON out of the content
    let parsed: any = null;
    try {
      parsed = JSON.parse(content);
    } catch (e) {
      // fallback: extract first JSON object
      const m = content.match(/(\{[\s\S]*\})/);
      if (m) {
        try {
          parsed = JSON.parse(m[1]);
        } catch (e2) {
          console.warn("verifyExecutionPlanWithLLM: failed to parse LLM response JSON");
        }
      }
    }

    if (parsed && typeof parsed === "object") {
      const approved = parsed.approved !== false; // default true if missing
      const suggested = parsed.executionPlan || {};
      // attach llmReview metadata
      const newPlan = { ...plan, ...suggested };
      newPlan._llmReview = { approved: !!approved, notes: parsed.notes || null };
      decision.executionPlan = newPlan as any;
    }

    return decision;
  } catch (err) {
    console.warn("verifyExecutionPlanWithLLM error:", err);
    return decision;
  }
}