import Groq from "groq-sdk";
import { zodToJsonSchema } from "zod-to-json-schema";
import type { LogLine, AvailableModel, CreateChatCompletionOptions, LLMResponse } from "@browserbasehq/stagehand";
import { LLMClient } from "@browserbasehq/stagehand";

type LLMCache = any;

export interface GroqClientOptions {
  apiKey?: string;
  baseURL?: string;
  maxRetries?: number;
  timeout?: number;
  skipTokenLimitCheck?: boolean;  // Skip token limit validation entirely
}

// Groq API has specific parameter limitations
const MIN_TEMPERATURE = 1e-8; // Minimum allowed temperature value

// List of unsupported OpenAI parameters that should be filtered out
const UNSUPPORTED_OPENAI_PARAMS = new Set([
  "logprobs",
  "logit_bias",
  "top_logprobs",
]);

// Default token limit for unknown models (conservative)
const DEFAULT_TOKEN_LIMIT = 8192;

// Model-specific token limits (context window sizes)
const MODEL_TOKEN_LIMITS: Record<string, number> = {
  // Production Models
  "mixtral-8x7b-32768": 32768,
  "llama-3.3-70b-versatile": 128000,
  "qwen-2.5-32b": 128000,
  // Preview Models
  "deepseek-r1-distill-llama-70b": 128000,
};

// Maximum output tokens per model (if different from context window)
const MODEL_MAX_OUTPUT_TOKENS: Record<string, number> = {
  "llama-3.3-70b-versatile": 32768,
  "llama-3.1-8b-instant": 8192,
  "qwen-2.5-32b": 8000,
  "deepseek-r1-distill-llama-70b": 8000,
};

// Groq API response types
export interface GroqMessage {
  role: "assistant" | "system" | "user" | "tool";
  content: string | null;
  tool_calls?: Array<GroqToolCall>;
}

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

export interface GroqCompletion {
  id: string;
  object: "chat.completion";
  created: number;
  model: string;
  choices: Array<{
    index: number;
    message: GroqMessage;
    finish_reason: "stop" | "length" | "tool_calls" | "content_filter" | null;
  }>;
  usage: {
    prompt_tokens: number;
    completion_tokens: number;
    total_tokens: number;
  };
}

// Groq-specific error classes
export class GroqError extends Error {
  constructor(message: string) {
    super(message);
    this.name = "GroqError";
  }
}

export class GroqAPIError extends GroqError {
  status: number;
  code: string;
  param: string | null;

  constructor(
    message: string,
    status: number,
    code: string,
    param: string | null = null,
  ) {
    super(message);
    this.name = "GroqAPIError";
    this.status = status;
    this.code = code;
    this.param = param;
  }
}

export class GroqAuthenticationError extends GroqAPIError {
  constructor(message: string) {
    super(message, 401, "invalid_api_key");
    this.name = "GroqAuthenticationError";
  }
}

export class GroqRateLimitError extends GroqAPIError {
  constructor(message: string) {
    super(message, 429, "rate_limit_exceeded");
    this.name = "GroqRateLimitError";
  }
}

export class GroqTimeoutError extends GroqError {
  constructor(message: string = "Request timed out") {
    super(message);
    this.name = "GroqTimeoutError";
  }
}

export class GroqConnectionError extends GroqError {
  constructor(message: string = "Failed to connect to Groq API") {
    super(message);
    this.name = "GroqConnectionError";
  }
}

export class GroqValidationError extends GroqError {
  constructor(message: string) {
    super(message);
    this.name = "GroqValidationError";
  }
}

/**
 * Determines if an error is retryable based on its type and message.
 * @param error - The error to check
 * @returns boolean indicating if the error is retryable
 */
function isRetryableError(error: any): boolean {
  return (
    error.message?.includes("Rate limit") ||
    error.error?.code === "rate_limit_exceeded" ||
    error.message?.includes("timeout") ||
    error.message?.includes("timed out") ||
    error.error?.code === "timeout"
  );
}

function isConnectionError(error: any): boolean {
  return (
    error.message?.includes("ECONNREFUSED") ||
    error.message?.includes("ECONNRESET") ||
    error.message?.includes("ETIMEDOUT") ||
    error.message?.includes("connect timeout")
  );
}

/**
 * Parses a failed tool call generation error.
 * @param error - The error containing failed generation data
 * @returns Parsed arguments from the failed generation
 * @throws GroqValidationError if parsing fails
 */
function parseFailedToolCallGeneration(error: any): any {
  try {
    // First attempt: direct JSON parse
    return JSON.parse(error.error.error.failed_generation).arguments;
  } catch (parseError) {
    try {
      // Second attempt: extract JSON from markdown code block
      const jsonMatch = error.error.error.failed_generation.match(/```json\n([\s\S]*?)\n```/);
      if (!jsonMatch) {
        throw new Error("No JSON code block found in failed generation");
      }
      return JSON.parse(jsonMatch[1].trim()).arguments;
    } catch (secondError) {
      throw new GroqValidationError(
        `Failed to parse tool call output: ${secondError instanceof Error ? secondError.message : String(secondError)}`
      );
    }
  }
}

export class GroqClient extends LLMClient {
  public type = "groq" as const;
  private client: Groq;
  // TODO: need to export the cache class from stagehand
  private cache: LLMCache | undefined;
  private enableCaching: boolean;
  public clientOptions: GroqClientOptions;
  public userProvidedInstructions?: string; // this is required by the LLMClient interface

  /**
   * Combines system prompts with user-provided instructions.
   * @param systemPrompt - The base system prompt
   * @param userProvidedInstructions - Optional user-provided instructions
   * @returns Combined system prompt with user instructions
   */
  private combineSystemPrompt(
    systemPrompt: string,
    userProvidedInstructions?: string,
  ): string {
    if (!userProvidedInstructions?.trim()) {
      return systemPrompt;
    }

    return `${systemPrompt.trim()}

# Custom Instructions Provided by the User

Please keep the user's instructions in mind when performing actions. If the user's instructions are not relevant to the current task, ignore them.

User Instructions:
${userProvidedInstructions}`;
  }

  /**
   * Prepares messages array with system prompt and user instructions.
   * @param messages - Original messages array
   * @param logger - Logger function
   * @returns Modified messages array with combined system prompt
   */
  private prepareMessages(
    messages: CreateChatCompletionOptions["options"]["messages"],
    logger: (message: LogLine) => void,
  ): CreateChatCompletionOptions["options"]["messages"] {
    // Find the first system message if it exists
    const systemMessageIndex = messages.findIndex(
      (msg) => msg.role === "system",
    );

    if (systemMessageIndex === -1 && !this.userProvidedInstructions?.trim()) {
      // No system message and no user instructions, return original messages
      return messages;
    }

    const newMessages = [...messages];
    if (systemMessageIndex === -1) {
      // No system message but we have user instructions, add a new system message
      logger({
        category: "groq",
        message: "Adding system message with user instructions",
        level: 1,
        auxiliary: {
          userInstructions: {
            value: this.userProvidedInstructions || "",
            type: "string",
          },
        },
      });

      newMessages.unshift({
        role: "system",
        content: this.combineSystemPrompt(
          "You are a helpful assistant.",
          this.userProvidedInstructions,
        ),
      });
    } else if (this.userProvidedInstructions?.trim()) {
      // Combine existing system message with user instructions
      const existingSystemMessage = newMessages[systemMessageIndex];
      const existingContent = typeof existingSystemMessage.content === "string"
        ? existingSystemMessage.content
        : JSON.stringify(existingSystemMessage.content);

      logger({
        category: "groq",
        message: "Combining system message with user instructions",
        level: 1,
        auxiliary: {
          originalSystemMessage: {
            value: existingContent,
            type: "string",
          },
          userInstructions: {
            value: this.userProvidedInstructions || "",
            type: "string",
          },
        },
      });

      newMessages[systemMessageIndex] = {
        role: "system",
        content: this.combineSystemPrompt(
          existingContent,
          this.userProvidedInstructions,
        ),
      };
    } else {
      // We have a system message but no user instructions
      // If the content is not a string, stringify it
      const existingSystemMessage = newMessages[systemMessageIndex];
      if (typeof existingSystemMessage.content !== "string") {
        newMessages[systemMessageIndex] = {
          role: "system",
          content: JSON.stringify(existingSystemMessage.content),
        };
      }
    }

    return newMessages;
  }

  private validateParameters(
    options: CreateChatCompletionOptions["options"],
    logger: (message: LogLine) => void,
  ) {
    // Check for unsupported parameters
    const unsupportedParams = Object.keys(options).filter((param) =>
      UNSUPPORTED_OPENAI_PARAMS.has(param),
    );

    if (unsupportedParams.length > 0) {
      logger({
        category: "groq",
        message: "Unsupported OpenAI parameters detected and will be ignored",
        level: 1,
        auxiliary: {
          unsupportedParams: {
            value: JSON.stringify(unsupportedParams),
            type: "object",
          },
        },
      });
    }

    // Validate messages array
    if (!options.messages.length) {
      throw new Error("Messages array cannot be empty");
    }

    // Validate message roles
    const validRoles = new Set(["system", "user", "assistant", "function"]);
    const invalidRole = options.messages.find(
      (msg) => !validRoles.has(msg.role),
    );
    if (invalidRole) {
      throw new Error(`Invalid message role: ${invalidRole.role}`);
    }

    // Skip token limit validation if configured
    if (this.clientOptions.skipTokenLimitCheck) {
      logger({
        category: "groq",
        message: "Skipping token limit validation as configured",
        level: 1,
        auxiliary: {
          model: {
            value: this.modelName,
            type: "string",
          },
          maxTokens: {
            value: String(options.maxTokens),
            type: "integer",
          },
        },
      });
      return;
    }

    // Validate max_tokens against model limits
    if (options.maxTokens) {
      const contextLimit = MODEL_TOKEN_LIMITS[this.modelName] ?? DEFAULT_TOKEN_LIMIT;
      const outputLimit = MODEL_MAX_OUTPUT_TOKENS[this.modelName] ?? contextLimit;
      const effectiveLimit = Math.min(contextLimit, outputLimit);

      // Log if we're using default limits for an unknown model
      if (!MODEL_TOKEN_LIMITS[this.modelName]) {
        logger({
          category: "groq",
          message: `Using default token limit (${DEFAULT_TOKEN_LIMIT}) for unknown model ${this.modelName}`,
          level: 1,
          auxiliary: {
            model: {
              value: this.modelName,
              type: "string",
            },
            defaultLimit: {
              value: String(DEFAULT_TOKEN_LIMIT),
              type: "integer",
            },
          },
        });
      }

      if (options.maxTokens > effectiveLimit) {
        logger({
          category: "groq",
          message: `max_tokens (${options.maxTokens}) exceeds model limit (${effectiveLimit}), will be capped`,
          level: 1,
          auxiliary: {
            maxTokens: {
              value: String(options.maxTokens),
              type: "integer",
            },
            modelLimit: {
              value: String(effectiveLimit),
              type: "integer",
            },
            contextWindow: {
              value: String(contextLimit),
              type: "integer",
            },
            maxOutputTokens: {
              value: String(outputLimit),
              type: "integer",
            },
          },
        });
        options.maxTokens = effectiveLimit;
      }
    }
  }

  constructor({
    enableCaching = false,
    cache,
    modelName,
    clientOptions,
    userProvidedInstructions,
  }: {
    enableCaching?: boolean;
    cache?: LLMCache;
    modelName: AvailableModel;
    clientOptions?: GroqClientOptions;
    userProvidedInstructions?: string;
  }) {
    super(modelName);
    this.client = new Groq({
      apiKey: clientOptions?.apiKey,
      baseURL: clientOptions?.baseURL,
      maxRetries: clientOptions?.maxRetries ?? 2,
      timeout: clientOptions?.timeout ?? 60000,
    });
    this.cache = cache;
    this.enableCaching = enableCaching;
    this.modelName = modelName;
    this.clientOptions = clientOptions || {};
    this.userProvidedInstructions = userProvidedInstructions;
  }

  async createChatCompletion<T = LLMResponse>({
    options,
    retries = 3,
    logger,
  }: CreateChatCompletionOptions): Promise<T> {
    logger({
      category: "groq",
      message: "Creating chat completion with Groq",
      level: 1,
      auxiliary: {
        options: {
          value: JSON.stringify(options),
          type: "object",
        },
        modelName: {
          value: this.modelName,
          type: "string",
        },
      },
    });

    try {
      // Validate parameters before proceeding
      this.validateParameters(options, logger);

      // Prepare messages with system prompt and user instructions
      const messages = this.prepareMessages(options.messages, logger);

      // Prepare cache options
      const cacheOptions: Record<string, unknown> = {
        model: this.modelName,
        messages: messages,
        temperature: options.temperature,
        top_p: options.top_p,
        frequency_penalty: options.frequency_penalty,
        presence_penalty: options.presence_penalty,
        maxTokens: options.maxTokens,
        tools: options.tools,
        tool_choice: options.tool_choice,
      };

      // Check cache if enabled
      if (this.enableCaching && this.cache) {
        //@ts-ignore
        const cachedResponse = await this.cache.get<T>(
          cacheOptions,
          options.requestId,
        );
        if (cachedResponse) {
          logger({
            category: "groq_cache",
            message: "Returning cached response",
            level: 1,
            auxiliary: {
              response: {
                value: JSON.stringify(cachedResponse),
                type: "object",
              },
            },
          });
          return cachedResponse;
        }
      }

      // Filter out messages with 'name' property as it's not supported by Groq
      const messagesFiltered = messages.map((msg) => ({
        role: msg.role,
        content: Array.isArray(msg.content)
          ? msg.content
              .map((c) => {
                if (typeof c === "string") return c;
                if ("text" in c) return c.text || "";
                if ("image_url" in c && c.image_url) return c.image_url.url;
                return "";
              })
              .join("\n")
          : msg.content,
      }));

      // Handle response model if provided
      let functionDefinition;
      if (options.response_model) {
        const jsonSchema = zodToJsonSchema(options.response_model.schema);
        functionDefinition = {
          name: "print_extracted_data",
          description: `Print the extracted data in the following format: ${options.response_model.name}`,
          parameters: jsonSchema,
        };
      }

      // Prepare tools if provided
      const tools =
        options.tools?.map((tool) => ({
          type: "function" as const,
          function: {
            name: tool.name,
            description: tool.description,
            parameters: tool.parameters,
          },
        })) || []

      if(functionDefinition) {
        tools.push({
          type: "function" as const,
          function: functionDefinition,
        });
      }

      try {
        const completion = await this.client.chat.completions.create({
          model: this.modelName,
          messages: messagesFiltered,
          temperature:
            options.temperature === 0 ? MIN_TEMPERATURE : options.temperature,
          top_p: options.top_p,
          frequency_penalty: options.frequency_penalty,
          presence_penalty: options.presence_penalty,
          max_tokens: options.maxTokens,
          ...(this.modelName.includes("r1") ? {reasoning_format: "hidden"} : {}), // TODO might be worth switching to "parsed" for transparency see https://console.groq.com/docs/reasoning
          n: 1, // Groq only supports n=1
          tools,
          tool_choice: options.response_model ? "required" : options.tool_choice as
            | "none"
            | "auto"
            | "required"
            | undefined,
        });

        // Transform to LLMResponse format
        const response: LLMResponse = {
          id: completion.id,
          object: "chat.completion",
          created: Date.now(),
          model: completion.model,
          choices: [
            {
              index: 0,
              message: {
                role: completion.choices[0].message.role,
                content: completion.choices[0].message.content || null,
                tool_calls: completion.choices[0].message.tool_calls || [],
              },
              finish_reason: completion.choices[0].finish_reason,
            },
          ],
          usage: {
            prompt_tokens: completion.usage?.prompt_tokens || 0,
            completion_tokens: completion.usage?.completion_tokens || 0,
            total_tokens: completion.usage?.total_tokens || 0,
          },
        };

        // Cache the response if caching is enabled
        if (this.enableCaching && this.cache) {
          await this.cache.set(cacheOptions, response, options.requestId);
        }

        // Handle response model extraction if needed
        if (
          options.response_model &&
          response.choices[0].message.tool_calls?.length > 0
        ) {
          const toolCall = response.choices[0].message.tool_calls[0];
          if (toolCall.function.name === "print_extracted_data") {
            try {
              return JSON.parse(toolCall.function.arguments) as T;
            } catch (error) {
              const validationError = new GroqValidationError(
                `Failed to parse structured output: ${error instanceof Error ? error.message : String(error)}`
              );
              logger({
                category: "groq",
                message: "Failed to parse structured output",
                level: 2,
                auxiliary: {
                  error: {
                    value:
                      error instanceof Error ? error.message : String(error),
                    type: "string",
                  },
                  toolCall: {
                    value: JSON.stringify(toolCall),
                    type: "object",
                  },
                },
              });
              if (retries > 0) {
                return this.createChatCompletion({
                  options,
                  retries: retries - 1,
                  logger,
                });
              }
              throw validationError;
            }
          }
        }

        return response as T;
      } catch (error) {
        const errorObj = error as any;
        
        // Re-throw GroqValidationError immediately
        if (errorObj instanceof GroqValidationError) {
          throw errorObj;
        }

        // Re-throw GroqError instances immediately
        if (errorObj instanceof GroqError) {
          throw errorObj;
        }

        // Handle retryable errors first
        if (isRetryableError(errorObj)) {
          if (retries > 0) {
            logger({
              category: "groq",
              message: "Retrying request due to retryable error",
              level: 1, // Info level
              auxiliary: {
                error: {
                  value: errorObj.message || String(errorObj),
                  type: "string",
                },
                retriesRemaining: {
                  value: String(retries - 1),
                  type: "integer",
                },
              },
            });
            return this.createChatCompletion({
              options,
              retries: retries - 1,
              logger,
            });
          }
          // Throw appropriate error type based on the error
          if (errorObj.message?.includes("timeout") || errorObj.message?.includes("timed out") || errorObj.error?.code === "timeout") {
            throw new GroqTimeoutError(errorObj.message);
          }
          throw new GroqRateLimitError(errorObj.message);
        }

        // Handle connection errors
        if (isConnectionError(errorObj)) {
          throw new GroqConnectionError(errorObj.message);
        }

        // Handle authentication errors
        if (errorObj.error?.code === "invalid_api_key" || errorObj.message?.includes("Invalid API Key")) {
          throw new GroqAuthenticationError(errorObj.message || "Invalid API key");
        }

        // Handle failed tool call generation
        if (errorObj.error?.error?.failed_generation !== undefined) {
          try {
            if (errorObj.error.error.failed_generation === null) {
              throw new GroqValidationError('Failed generation was null');
            }
            return parseFailedToolCallGeneration(errorObj);
          } catch (parseError) {
            if (parseError instanceof GroqValidationError) {
              throw parseError;
            }
            throw new GroqValidationError(
              `Failed to parse tool call output: ${parseError instanceof Error ? parseError.message : String(parseError)}`
            );
          }
        }

        // If it's a generic error without API-specific properties, throw a generic error
        if (!errorObj.error && !errorObj.status && !errorObj.code) {
          throw new Error('Unhandled Groq error');
        }

        // Handle API errors with status codes
        const status = errorObj.status || errorObj.error?.status || 500;
        const code = errorObj.code || errorObj.error?.code || null;
        const message = errorObj.error?.message || errorObj.message || String(errorObj);
        const param = errorObj.error?.param || errorObj.param || null;
        throw new GroqAPIError(
          `Groq API error (${status}): ${message}`,
          status,
          code,
          param
        );
      }
    } catch (error) {
      // Log the error through the logger
      logger({
        category: "groq",
        message: "Unhandled error in Groq chat completion",
        level: 2, // Error level
        auxiliary: {
          error: {
            value: error instanceof Error ? error.message : String(error),
            type: "string",
          },
          trace: {
            value: error instanceof Error && error.stack ? error.stack : "",
            type: "string",
          },
          errorObject: {
            value: JSON.stringify(error, null, 2),
            type: "object",
          },
        },
      });

      // Re-throw the error if it's already one of our custom error types
      if (error instanceof GroqError) {
        throw error;
      }

      // Otherwise, create a more informative generic error
      throw new GroqAPIError(
        `Unhandled Groq error: ${error instanceof Error ? error.message : String(error)}`,
        500,
        "internal_error",
        null
      );
    }
  }
}