import { TOOL_DEFINITIONS, executeTool, type ToolContext } from "./db-tools";

const LITELLM_URL = process.env.LITELLM_BASE_URL;
const LITELLM_KEY = process.env.LITELLM_API_KEY;
const QUERY_MODEL = process.env.QUERY_MODEL || "minimax-2.7";

interface Message {
  role: "system" | "user" | "assistant" | "tool";
  content: string;
  tool_calls?: ToolCall[];
  tool_call_id?: string;
  name?: string;
}

interface ToolCall {
  id: string;
  type: "function";
  function: { name: string; arguments: string };
}

export async function answerStructuredQuery(query: string, ctx: ToolContext): Promise<string> {
  if (!LITELLM_URL || !LITELLM_KEY) return "AI not configured.";

  const now = new Date();
  const systemPrompt = `You are a data assistant for ${ctx.childName}'s baby tracking app. Today is ${now.toLocaleDateString("en-IN", { weekday: "long", day: "numeric", month: "long" })}. Use the available tools to fetch data and answer the parent's question with specific numbers and times. Keep responses concise and warm. Use "she/her" or "he/his" based on context if unknown use their name. Never interpret symptoms or give medical advice.`;

  const messages: Message[] = [
    { role: "system", content: systemPrompt },
    { role: "user", content: query },
  ];

  // Tool-use loop (max 3 iterations)
  for (let i = 0; i < 3; i++) {
    const res = await fetch(`${LITELLM_URL}/v1/chat/completions`, {
      method: "POST",
      headers: { "Content-Type": "application/json", "Authorization": `Bearer ${LITELLM_KEY}` },
      body: JSON.stringify({
        model: QUERY_MODEL,
        messages,
        tools: TOOL_DEFINITIONS,
        tool_choice: "auto",
        temperature: 0,
        max_tokens: 400,
      }),
    });

    if (!res.ok) {
      const err = await res.text();
      console.error("[structured-query] LLM error:", err);
      break;
    }

    const data = await res.json();
    const choice = data.choices?.[0];
    const msg = choice?.message;

    if (!msg) break;

    // If the model wants to call tools
    if (msg.tool_calls?.length) {
      messages.push({ role: "assistant", content: msg.content || "", tool_calls: msg.tool_calls });

      // Execute each tool call in parallel
      const toolResults = await Promise.all(
        msg.tool_calls.map(async (tc: ToolCall) => {
          let args: Record<string, unknown> = {};
          try { args = JSON.parse(tc.function.arguments); } catch { /* empty args */ }
          const result = await executeTool(tc.function.name, args, ctx).catch(e => ({ error: String(e) }));
          return {
            role: "tool" as const,
            tool_call_id: tc.id,
            name: tc.function.name,
            content: JSON.stringify(result),
          };
        })
      );

      messages.push(...toolResults);
      continue;
    }

    // Final text response
    if (choice?.finish_reason === "stop" && msg.content) {
      return msg.content;
    }

    break;
  }

  return "I couldn't retrieve that data right now. Please try again.";
}