import OpenAI from "openai";
import { zodResponseFormat } from "openai/helpers/zod";
import { z } from "zod";
import { LLM, LLMResponse } from "./base";
import { LLMConfig, Message } from "../types";

/**
 * OpenAI LLM implementation with structured output support.
 * This is the default implementation for the "openai" provider.
 */
export class OpenAIStructuredLLM implements LLM {
  private openai: OpenAI;
  private model: string;

  constructor(config: LLMConfig) {
    this.openai = new OpenAI({ apiKey: config.apiKey });
    this.model = config.model || "gpt-4o-2024-08-06"; // Use model that supports structured output
  }

  async generateResponse(
    messages: Message[],
    responseFormat?: { type: string } | null,
    tools?: any[],
    zodSchema?: z.ZodType<any> // Add support for zod schemas
  ): Promise<string | LLMResponse> {
    // If we have a zod schema, try structured output first
    if (zodSchema) {
      try {
        // Ensure we're using a model that supports structured output
        const structuredModel = this.model.includes('gpt-4o') ? this.model : 'gpt-4o-2024-08-06';
        
        const completion = await this.openai.beta.chat.completions.parse({
          messages: messages.map((msg) => ({
            role: msg.role as "system" | "user" | "assistant",
            content:
              typeof msg.content === "string"
                ? msg.content
                : JSON.stringify(msg.content),
          })),
          model: structuredModel,
          response_format: zodResponseFormat(zodSchema, "structured_response"),
        });

        const response = completion.choices[0].message;
        
        // Return the parsed structured data
        return {
          content: JSON.stringify(response.parsed),
          role: response.role,
          parsed: response.parsed, // Include the parsed object
        };
      } catch (error) {
        console.warn("Structured output failed, falling back to JSON mode:", (error as Error).message);
        
        // Fallback to JSON mode with the original model
        try {
          const completion = await this.openai.chat.completions.create({
            messages: messages.map((msg) => ({
              role: msg.role as "system" | "user" | "assistant",
              content:
                typeof msg.content === "string"
                  ? msg.content + "\n\nPlease respond with valid JSON matching the expected schema."
                  : JSON.stringify(msg.content),
            })),
            model: this.model,
            response_format: { type: "json_object" },
          });

          const response = completion.choices[0].message;
          
          // Try to parse the JSON response
          try {
            const parsed = JSON.parse(response.content || "{}");
            return {
              content: response.content || "",
              role: response.role,
              parsed: parsed,
            };
          } catch (parseError) {
            console.warn("Failed to parse JSON fallback response");
            return response.content || "";
          }
        } catch (fallbackError) {
          console.warn("JSON mode fallback also failed:", (fallbackError as Error).message);
          // Final fallback to plain text
          return await this.generatePlainResponse(messages);
        }
      }
    }

    // Original implementation for non-structured responses
    const completion = await this.openai.chat.completions.create({
      messages: messages.map((msg) => ({
        role: msg.role as "system" | "user" | "assistant",
        content:
          typeof msg.content === "string"
            ? msg.content
            : JSON.stringify(msg.content),
      })),
      model: this.model,
      ...(tools
        ? {
            tools: tools.map((tool) => ({
              type: "function",
              function: {
                name: tool.function.name,
                description: tool.function.description,
                parameters: tool.function.parameters,
              },
            })),
            tool_choice: "auto" as const,
          }
        : responseFormat
          ? {
              response_format: {
                type: responseFormat.type as "text" | "json_object",
              },
            }
          : {}),
    });

    const response = completion.choices[0].message;

    if (response.tool_calls) {
      return {
        content: response.content || "",
        role: response.role,
        toolCalls: response.tool_calls.map((call) => ({
          name: call.function.name,
          arguments: call.function.arguments,
        })),
      };
    }

    return response.content || "";
  }

  private async generatePlainResponse(messages: Message[]): Promise<string> {
    const completion = await this.openai.chat.completions.create({
      messages: messages.map((msg) => ({
        role: msg.role as "system" | "user" | "assistant",
        content:
          typeof msg.content === "string"
            ? msg.content
            : JSON.stringify(msg.content),
      })),
      model: this.model,
    });

    return completion.choices[0].message.content || "";
  }

  async generateChat(messages: Message[]): Promise<LLMResponse> {
    const completion = await this.openai.chat.completions.create({
      messages: messages.map((msg) => ({
        role: msg.role as "system" | "user" | "assistant",
        content:
          typeof msg.content === "string"
            ? msg.content
            : JSON.stringify(msg.content),
      })),
      model: this.model,
    });
    const response = completion.choices[0].message;
    return {
      content: response.content || "",
      role: response.role,
    };
  }
}