/**
 * Antigravity API Integration for Binary-Free Mode
 * Provides direct access to Antigravity models (Gemini, Claude, GPT-OSS) via Google OAuth
 * Based on the opencode-antigravity-auth plugin: https://github.com/NoeFabris/opencode-antigravity-auth
 * 
 * This integration enables access to:
 * - Gemini 3 Pro/Flash models
 * - Claude Sonnet 4.5 / Opus 4.5 (with thinking variants)
 * - GPT-OSS 120B
 * 
 * Uses Google OAuth credentials stored via the Antigravity OAuth flow
 */

import { z } from "zod"

// Configuration schema for Antigravity
export const AntigravityConfigSchema = z.object({
    enabled: z.boolean().default(true),
    // Multiple endpoints for automatic fallback (daily → autopush → prod)
    endpoints: z.array(z.string()).default([
        "https://daily.antigravity.dev/v1beta",
        "https://autopush.antigravity.dev/v1beta",
        "https://antigravity.dev/v1beta"
    ]),
    apiKey: z.string().optional()
})

export type AntigravityConfig = z.infer<typeof AntigravityConfigSchema>

// Antigravity Model schema
export const AntigravityModelSchema = z.object({
    id: z.string(),
    name: z.string(),
    family: z.enum(["gemini", "claude", "gpt-oss"]).optional(),
    reasoning: z.boolean().optional(),
    tool_call: z.boolean().optional(),
    limit: z.object({
        context: z.number(),
        output: z.number()
    }).optional()
})

export type AntigravityModel = z.infer<typeof AntigravityModelSchema>

// Chat message schema (OpenAI-compatible)
export const ChatMessageSchema = z.object({
    role: z.enum(["user", "assistant", "system", "tool"]),
    content: z.string().optional(),
    tool_calls: z.array(z.object({
        id: z.string(),
        type: z.literal("function"),
        function: z.object({
            name: z.string(),
            arguments: z.string()
        })
    })).optional(),
    tool_call_id: z.string().optional()
})

export type ChatMessage = z.infer<typeof ChatMessageSchema>

// Tool Definition Schema
export const ToolDefinitionSchema = z.object({
    type: z.literal("function"),
    function: z.object({
        name: z.string(),
        description: z.string(),
        parameters: z.object({
            type: z.literal("object"),
            properties: z.record(z.any()),
            required: z.array(z.string()).optional()
        })
    })
})

export type ToolDefinition = z.infer<typeof ToolDefinitionSchema>

// Chat request schema
export const ChatRequestSchema = z.object({
    model: z.string(),
    messages: z.array(ChatMessageSchema),
    stream: z.boolean().default(true),
    temperature: z.number().optional(),
    max_tokens: z.number().optional(),
    tools: z.array(ToolDefinitionSchema).optional(),
    tool_choice: z.union([
        z.literal("auto"),
        z.literal("none"),
        z.object({
            type: z.literal("function"),
            function: z.object({ name: z.string() })
        })
    ]).optional(),
    workspacePath: z.string().optional(),
    enableTools: z.boolean().optional()
})

export type ChatRequest = z.infer<typeof ChatRequestSchema>

// Chat response chunk schema
export const ChatChunkSchema = z.object({
    id: z.string().optional(),
    object: z.string().optional(),
    created: z.number().optional(),
    model: z.string().optional(),
    choices: z.array(z.object({
        index: z.number(),
        delta: z.object({
            role: z.string().optional(),
            content: z.string().optional(),
            reasoning_content: z.string().optional()
        }).optional(),
        message: z.object({
            role: z.string(),
            content: z.string()
        }).optional(),
        finish_reason: z.string().nullable().optional()
    }))
})

export type ChatChunk = z.infer<typeof ChatChunkSchema>

// Available Antigravity models with their specifications
export const ANTIGRAVITY_MODELS: AntigravityModel[] = [
    // Gemini Models
    {
        id: "gemini-3-pro-low",
        name: "Gemini 3 Pro Low (Antigravity)",
        family: "gemini",
        reasoning: true,
        tool_call: true,
        limit: { context: 1048576, output: 65535 }
    },
    {
        id: "gemini-3-pro-high",
        name: "Gemini 3 Pro High (Antigravity)",
        family: "gemini",
        reasoning: true,
        tool_call: true,
        limit: { context: 1048576, output: 65535 }
    },
    {
        id: "gemini-3-flash",
        name: "Gemini 3 Flash (Antigravity)",
        family: "gemini",
        reasoning: false,
        tool_call: true,
        limit: { context: 1048576, output: 65536 }
    },
    // Claude Models
    {
        id: "claude-sonnet-4-5",
        name: "Claude Sonnet 4.5 (Antigravity)",
        family: "claude",
        reasoning: false,
        tool_call: true,
        limit: { context: 200000, output: 64000 }
    },
    {
        id: "claude-sonnet-4-5-thinking-low",
        name: "Claude Sonnet 4.5 Thinking Low (Antigravity)",
        family: "claude",
        reasoning: true,
        tool_call: true,
        limit: { context: 200000, output: 64000 }
    },
    {
        id: "claude-sonnet-4-5-thinking-medium",
        name: "Claude Sonnet 4.5 Thinking Medium (Antigravity)",
        family: "claude",
        reasoning: true,
        tool_call: true,
        limit: { context: 200000, output: 64000 }
    },
    {
        id: "claude-sonnet-4-5-thinking-high",
        name: "Claude Sonnet 4.5 Thinking High (Antigravity)",
        family: "claude",
        reasoning: true,
        tool_call: true,
        limit: { context: 200000, output: 64000 }
    },
    {
        id: "claude-opus-4-5-thinking-low",
        name: "Claude Opus 4.5 Thinking Low (Antigravity)",
        family: "claude",
        reasoning: true,
        tool_call: true,
        limit: { context: 200000, output: 64000 }
    },
    {
        id: "claude-opus-4-5-thinking-medium",
        name: "Claude Opus 4.5 Thinking Medium (Antigravity)",
        family: "claude",
        reasoning: true,
        tool_call: true,
        limit: { context: 200000, output: 64000 }
    },
    {
        id: "claude-opus-4-5-thinking-high",
        name: "Claude Opus 4.5 Thinking High (Antigravity)",
        family: "claude",
        reasoning: true,
        tool_call: true,
        limit: { context: 200000, output: 64000 }
    },
    // Other Models
    {
        id: "gpt-oss-120b-medium",
        name: "GPT-OSS 120B Medium (Antigravity)",
        family: "gpt-oss",
        reasoning: true,
        tool_call: true,
        limit: { context: 131072, output: 32768 }
    }
]

// Token storage key for Antigravity OAuth
const ANTIGRAVITY_TOKEN_KEY = "antigravity_oauth_token"

export interface AntigravityToken {
    access_token: string
    refresh_token?: string
    expires_in: number
    created_at: number
    project_id?: string
}

export class AntigravityClient {
    private config: AntigravityConfig
    private currentEndpointIndex: number = 0
    private modelsCache: AntigravityModel[] | null = null
    private modelsCacheTime: number = 0
    private readonly CACHE_TTL_MS = 5 * 60 * 1000 // 5 minutes

    constructor(config?: Partial<AntigravityConfig>) {
        this.config = AntigravityConfigSchema.parse(config || {})
    }

    /**
     * Get the current active endpoint with automatic fallback
     */
    private getEndpoint(): string {
        const endpoints = this.config.endpoints
        return endpoints[this.currentEndpointIndex] || endpoints[0]
    }

    /**
     * Rotate to next endpoint on failure
     */
    private rotateEndpoint(): void {
        this.currentEndpointIndex = (this.currentEndpointIndex + 1) % this.config.endpoints.length
        console.log(`Antigravity: Rotating to endpoint ${this.getEndpoint()}`)
    }

    /**
     * Get stored OAuth token from localStorage (browser context)
     */
    getStoredToken(): AntigravityToken | null {
        if (typeof window === "undefined") return null
        try {
            const raw = window.localStorage.getItem(ANTIGRAVITY_TOKEN_KEY)
            if (!raw) return null
            return JSON.parse(raw)
        } catch {
            return null
        }
    }

    /**
     * Check if the stored token is still valid
     */
    isTokenValid(token: AntigravityToken | null): boolean {
        if (!token) return false
        const createdAt = token.created_at > 1e12 ? Math.floor(token.created_at / 1000) : token.created_at
        const expiresAt = (createdAt + token.expires_in) * 1000 - 300000 // 5 min buffer
        return Date.now() < expiresAt
    }

    /**
     * Get authorization headers for API requests
     */
    private getAuthHeaders(): Record<string, string> {
        const headers: Record<string, string> = {
            "Content-Type": "application/json",
            "User-Agent": "NomadArch/1.0"
        }

        // Try OAuth token first
        const token = this.getStoredToken()
        if (token && this.isTokenValid(token)) {
            headers["Authorization"] = `Bearer ${token.access_token}`
        } else if (this.config.apiKey) {
            headers["Authorization"] = `Bearer ${this.config.apiKey}`
        }

        return headers
    }

    /**
     * Check if the client is authenticated
     */
    isAuthenticated(): boolean {
        const token = this.getStoredToken()
        return this.isTokenValid(token) || Boolean(this.config.apiKey)
    }

    /**
     * Get available Antigravity models
     */
    async getModels(): Promise<AntigravityModel[]> {
        // Return cached models if still valid
        const now = Date.now()
        if (this.modelsCache && (now - this.modelsCacheTime) < this.CACHE_TTL_MS) {
            return this.modelsCache
        }

        // If authenticated, return full model list
        if (this.isAuthenticated()) {
            this.modelsCache = ANTIGRAVITY_MODELS
            this.modelsCacheTime = now
            return ANTIGRAVITY_MODELS
        }

        // Not authenticated - return empty list
        return []
    }

    /**
     * Test connection to Antigravity API
     */
    async testConnection(): Promise<boolean> {
        if (!this.isAuthenticated()) {
            return false
        }

        try {
            // Try a simple models list request to verify connectivity
            const response = await fetch(`${this.getEndpoint()}/models`, {
                headers: this.getAuthHeaders(),
                signal: AbortSignal.timeout(10000)
            })
            return response.ok
        } catch (error) {
            console.warn("Antigravity connection test failed:", error)
            // Try next endpoint
            this.rotateEndpoint()
            return false
        }
    }

    /**
     * Chat completion (streaming) with automatic endpoint fallback
     */
    async *chatStream(request: ChatRequest): AsyncGenerator<ChatChunk> {
        if (!this.isAuthenticated()) {
            throw new Error("Antigravity: Not authenticated. Please sign in with Google OAuth.")
        }

        let lastError: Error | null = null
        const maxRetries = this.config.endpoints.length

        for (let retry = 0; retry < maxRetries; retry++) {
            try {
                const endpoint = this.getEndpoint()
                const response = await fetch(`${endpoint}/chat/completions`, {
                    method: "POST",
                    headers: this.getAuthHeaders(),
                    body: JSON.stringify({
                        ...request,
                        stream: true
                    })
                })

                if (!response.ok) {
                    const errorText = await response.text()
                    if (response.status === 401 || response.status === 403) {
                        throw new Error(`Antigravity authentication failed: ${errorText}`)
                    }
                    // Try next endpoint for other errors
                    this.rotateEndpoint()
                    lastError = new Error(`Antigravity API error (${response.status}): ${errorText}`)
                    continue
                }

                if (!response.body) {
                    throw new Error("Response body is missing")
                }

                const reader = response.body.getReader()
                const decoder = new TextDecoder()
                let buffer = ""

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

                        buffer += decoder.decode(value, { stream: true })
                        const lines = buffer.split("\n")
                        buffer = lines.pop() || ""

                        for (const line of lines) {
                            const trimmed = line.trim()
                            if (trimmed.startsWith("data: ")) {
                                const data = trimmed.slice(6)
                                if (data === "[DONE]") return

                                try {
                                    const parsed = JSON.parse(data)
                                    yield parsed as ChatChunk

                                    // Check for finish
                                    if (parsed.choices?.[0]?.finish_reason) {
                                        return
                                    }
                                } catch (e) {
                                    // Skip invalid JSON
                                }
                            }
                        }
                    }
                } finally {
                    reader.releaseLock()
                }
                return // Success, exit retry loop
            } catch (error) {
                lastError = error instanceof Error ? error : new Error(String(error))
                if (error instanceof Error && error.message.includes("authentication")) {
                    throw error // Don't retry auth errors
                }
                this.rotateEndpoint()
            }
        }

        throw lastError || new Error("Antigravity: All endpoints failed")
    }

    /**
     * Chat completion (non-streaming)
     */
    async chat(request: ChatRequest): Promise<ChatChunk> {
        if (!this.isAuthenticated()) {
            throw new Error("Antigravity: Not authenticated. Please sign in with Google OAuth.")
        }

        let lastError: Error | null = null
        const maxRetries = this.config.endpoints.length

        for (let retry = 0; retry < maxRetries; retry++) {
            try {
                const endpoint = this.getEndpoint()
                const response = await fetch(`${endpoint}/chat/completions`, {
                    method: "POST",
                    headers: this.getAuthHeaders(),
                    body: JSON.stringify({
                        ...request,
                        stream: false
                    })
                })

                if (!response.ok) {
                    const errorText = await response.text()
                    if (response.status === 401 || response.status === 403) {
                        throw new Error(`Antigravity authentication failed: ${errorText}`)
                    }
                    this.rotateEndpoint()
                    lastError = new Error(`Antigravity API error (${response.status}): ${errorText}`)
                    continue
                }

                return await response.json()
            } catch (error) {
                lastError = error instanceof Error ? error : new Error(String(error))
                if (error instanceof Error && error.message.includes("authentication")) {
                    throw error
                }
                this.rotateEndpoint()
            }
        }

        throw lastError || new Error("Antigravity: All endpoints failed")
    }
}

export function getDefaultAntigravityConfig(): AntigravityConfig {
    return {
        enabled: true,
        endpoints: [
            "https://daily.antigravity.dev/v1beta",
            "https://autopush.antigravity.dev/v1beta",
            "https://antigravity.dev/v1beta"
        ]
    }
}
