interface Message {
  role: "system" | "user" | "assistant";
  content: string;
}

interface CompletionOptions {
  genstackToken: string;
  messages: Message[];
  model: string;
  developerId?: string;
  applicationId?: string;
  temperature?: number;
  maxOutputTokens?: number;
  overrideGrossMarginPercent?: number;
}

interface UsageInfo {
  type: "completion";
  promptTokens: number;
  completionTokens: number;
  inputCost: string;
  outputCost: string;
}

interface CompletionResponse {
  model: string;
  output: string;
  usage: UsageInfo;
  balanceBefore: string;
  balanceAfter: string;
}

class CompletionClient {
  private baseUrl: string;

  constructor(baseUrl: string) {
    this.baseUrl = baseUrl;
  }

  async generateCompletion(options: CompletionOptions): Promise<CompletionResponse> {
    const response = await fetch(this.baseUrl, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
      },
      body: JSON.stringify(options),
    });

    if (!response.ok) {
      const error = await response.json();
      throw new Error(error.error.message);
    }

    return response.json();
  }

  async generateStreamingCompletion(
    options: CompletionOptions,
    onToken: (token: string) => void,
    onComplete: (response: CompletionResponse) => void,
    onAbort?: (response: CompletionResponse) => void,
  ): Promise<string> {
    const response = await fetch(`${this.baseUrl}/stream`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
      },
      body: JSON.stringify(options),
    });

    if (!response.ok) {
      const error = await response.json();
      throw new Error(error.error.message);
    }

    const streamId = response.headers.get("X-Stream-Id");
    if (!streamId) {
      throw new Error("No stream ID received from server");
    }

    const reader = response.body!.getReader();
    const decoder = new TextDecoder();

    try {
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        const text = decoder.decode(value);
        const lines = text.split("\n");

        for (const line of lines) {
          if (!line.trim() || !line.startsWith("data: ")) continue;

          const data = JSON.parse(line.slice(6));
          if (data.token) {
            onToken(data.token);
          }
          if (data.completed) {
            onComplete(data.response);
            return streamId;
          }
          if (data.aborted) {
            onAbort?.(data.response);
            return streamId;
          }
        }
      }
    } catch (error) {
      reader.cancel();
      throw error;
    }

    return streamId;
  }

  async shutdownStream(streamId: string): Promise<void> {
    const response = await fetch(`${this.baseUrl}/stream/shutdown`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
      },
      body: JSON.stringify({ streamId }),
    });

    if (!response.ok) {
      const error = await response.json();
      throw new Error(error.error.message);
    }
  }

  async checkHealth(): Promise<{ status: string; env: string; ui: string }> {
    const response = await fetch(this.baseUrl);

    if (!response.ok) {
      throw new Error("Service health check failed");
    }

    return response.json();
  }
}

// Usage Examples
async function completionExamples() {
  const client = new CompletionClient("https://completion.genstack.app");

  try {
    // Basic completion
    const basicResult = await client.generateCompletion({
      genstackToken: "your_token_here",
      messages: [{ role: "user", content: "Tell me about AI." }],
      model: "gpt-3.5-turbo",
    });
    console.log("Completion:", basicResult.output);
    console.log("Token usage:", basicResult.usage);

    // Advanced completion
    const advancedResult = await client.generateCompletion({
      genstackToken: "your_token_here",
      messages: [
        { role: "system", content: "You are a helpful assistant." },
        { role: "user", content: "Tell me about AI." },
      ],
      model: "gpt-3.5-turbo",
      temperature: 0.7,
      maxOutputTokens: 500,
      developerId: "dev_123",
      applicationId: "app_456",
    });
    console.log("Advanced completion result:", advancedResult);
  } catch (error) {
    console.error("Completion failed:", error);
  }
}

async function handleCompletion(options: CompletionOptions) {
  try {
    const response = await fetch("https://completion.genstack.app/", {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
      },
      body: JSON.stringify(options),
    });

    if (!response.ok) {
      const error = await response.json();
      switch (error.error.type) {
        case "InsufficientFunds":
          console.error("Insufficient balance for completion");
          // Handle insufficient funds
          break;
        case "InvalidCredentials":
          console.error("Invalid or expired token");
          // Handle authentication error
          break;
        case "GatewayError":
          console.error("AI service error:", error.error.details);
          // Handle AI service errors
          break;
        default:
          console.error("Unexpected error:", error.error.message);
      }
      return;
    }

    const result = await response.json();
    // Process successful completion...
  } catch (error) {
    console.error("Request failed:", error);
  }
}