/**
 * 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 { randomUUID } from "crypto"
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-cloudcode-pa.sandbox.googleapis.com",
        "https://autopush-cloudcode-pa.sandbox.googleapis.com",
        "https://cloudcode-pa.googleapis.com"
    ]),
    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",
        name: "Claude Opus 4.5 (Antigravity)",
        family: "claude",
        reasoning: false,
        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"
const ANTIGRAVITY_HEADERS = {
    "User-Agent": "antigravity/1.11.5 windows/amd64",
    "X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1",
    "Client-Metadata": "{\"ideType\":\"IDE_UNSPECIFIED\",\"platform\":\"PLATFORM_UNSPECIFIED\",\"pluginType\":\"GEMINI\"}",
} as const
const LOAD_ASSIST_HEADERS = {
    "User-Agent": "google-api-nodejs-client/9.15.1",
    "X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1",
    "Client-Metadata": ANTIGRAVITY_HEADERS["Client-Metadata"],
} as const

const DEFAULT_PROJECT_ID = process.env.ANTIGRAVITY_PROJECT_ID || "rising-fact-p41fc"
const LOAD_ASSIST_METADATA = {
    ideType: "IDE_UNSPECIFIED",
    platform: "PLATFORM_UNSPECIFIED",
    pluginType: "GEMINI"
} as const
const LOAD_ENDPOINTS = [
    "https://cloudcode-pa.googleapis.com",
    "https://daily-cloudcode-pa.sandbox.googleapis.com",
    "https://autopush-cloudcode-pa.sandbox.googleapis.com"
] as const

const STREAM_ACTION = "streamGenerateContent"
const GENERATE_ACTION = "generateContent"

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

function generateSyntheticProjectId(): string {
    const adjectives = ["useful", "bright", "swift", "calm", "bold"]
    const nouns = ["fuze", "wave", "spark", "flow", "core"]
    const adj = adjectives[Math.floor(Math.random() * adjectives.length)]
    const noun = nouns[Math.floor(Math.random() * nouns.length)]
    const random = randomUUID().slice(0, 5).toLowerCase()
    return `${adj}-${noun}-${random}`
}

function collectSystemInstruction(messages: ChatMessage[]): string | undefined {
    const systemParts: string[] = []
    for (const message of messages) {
        if (message.role === "system" && typeof message.content === "string") {
            systemParts.push(message.content)
        }
    }
    const combined = systemParts.join("\n\n").trim()
    return combined.length > 0 ? combined : undefined
}

function buildContents(messages: ChatMessage[]): Array<{ role: "user" | "model"; parts: Array<{ text: string }> }> {
    const contents: Array<{ role: "user" | "model"; parts: Array<{ text: string }> }> = []
    for (const message of messages) {
        if (!message.content) continue
        if (message.role === "system") continue
        const role = message.role === "assistant" ? "model" : "user"
        const prefix = message.role === "tool" ? "Tool result:\n" : ""
        contents.push({
            role,
            parts: [{ text: `${prefix}${message.content}` }]
        })
    }
    return contents
}

function extractTextFromResponse(payload: any): string {
    const candidates = payload?.candidates
    if (!Array.isArray(candidates) || candidates.length === 0) return ""
    const parts = candidates[0]?.content?.parts
    if (!Array.isArray(parts)) return ""
    return parts.map((part: any) => (typeof part?.text === "string" ? part.text : "")).join("")
}

export class AntigravityClient {
    private config: AntigravityConfig
    private currentEndpointIndex: number = 0
    private modelsCache: AntigravityModel[] | null = null
    private modelsCacheTime: number = 0
    private projectIdCache: string | null = null
    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(accessToken?: string): Record<string, string> {
        const headers: Record<string, string> = {
            "Content-Type": "application/json",
            "User-Agent": ANTIGRAVITY_HEADERS["User-Agent"],
            "X-Goog-Api-Client": ANTIGRAVITY_HEADERS["X-Goog-Api-Client"],
            "Client-Metadata": ANTIGRAVITY_HEADERS["Client-Metadata"],
        }

        if (accessToken) {
            headers["Authorization"] = `Bearer ${accessToken}`
            return headers
        }

        // 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
    }

    private getLoadHeaders(accessToken: string): Record<string, string> {
        return {
            "Content-Type": "application/json",
            "User-Agent": LOAD_ASSIST_HEADERS["User-Agent"],
            "X-Goog-Api-Client": LOAD_ASSIST_HEADERS["X-Goog-Api-Client"],
            "Client-Metadata": LOAD_ASSIST_HEADERS["Client-Metadata"],
            Authorization: `Bearer ${accessToken}`,
        }
    }

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

    private async resolveProjectId(accessToken: string | undefined, projectIdOverride?: string): Promise<string> {
        const requestedProjectId = projectIdOverride?.trim()
        if (this.projectIdCache && !requestedProjectId) return this.projectIdCache
        if (!accessToken) {
            const fallback = requestedProjectId || generateSyntheticProjectId()
            if (requestedProjectId) {
                this.projectIdCache = requestedProjectId
            }
            return fallback
        }

        const loadEndpoints = Array.from(new Set([...LOAD_ENDPOINTS, ...this.config.endpoints]))
        const tryLoad = async (metadata: Record<string, string>): Promise<string | null> => {
            for (const endpoint of loadEndpoints) {
                try {
                    const response = await fetch(`${endpoint}/v1internal:loadCodeAssist`, {
                        method: "POST",
                        headers: this.getLoadHeaders(accessToken),
                        body: JSON.stringify({ metadata }),
                        signal: AbortSignal.timeout(10000),
                    })
                    if (!response.ok) continue
                    const data = await response.json() as any
                    const projectId =
                        data?.cloudaicompanionProject?.id ||
                        data?.cloudaicompanionProject ||
                        data?.projectId
                    if (typeof projectId === "string" && projectId.length > 0) {
                        return projectId
                    }
                } catch {
                    continue
                }
            }
            return null
        }

        let resolvedProjectId: string | null = null
        const baseMetadata: Record<string, string> = { ...LOAD_ASSIST_METADATA }
        if (requestedProjectId) {
            baseMetadata.duetProject = requestedProjectId
            resolvedProjectId = await tryLoad(baseMetadata)
        } else {
            resolvedProjectId = await tryLoad(baseMetadata)
            if (!resolvedProjectId) {
                const fallbackMetadata: Record<string, string> = {
                    ...LOAD_ASSIST_METADATA,
                    duetProject: DEFAULT_PROJECT_ID,
                }
                resolvedProjectId = await tryLoad(fallbackMetadata)
            }
        }

        const fallbackProjectId = requestedProjectId || DEFAULT_PROJECT_ID
        const finalProjectId = resolvedProjectId || fallbackProjectId
        this.projectIdCache = finalProjectId
        return finalProjectId
    }

    private resolveAccessToken(accessToken?: string): string | null {
        if (accessToken) return accessToken
        const token = this.getStoredToken()
        if (token && this.isTokenValid(token)) {
            return token.access_token
        }
        if (this.config.apiKey) {
            return this.config.apiKey
        }
        return null
    }

    private async requestGenerateContent(request: ChatRequest, accessToken?: string, projectIdOverride?: string): Promise<string> {
        const authToken = this.resolveAccessToken(accessToken)
        if (!authToken) {
            throw new Error("Antigravity: Missing access token.")
        }

        const projectId = await this.resolveProjectId(authToken, projectIdOverride)
        const systemInstruction = collectSystemInstruction(request.messages)
        const contents = buildContents(request.messages)

        const generationConfig: Record<string, unknown> = {}
        if (typeof request.temperature === "number") {
            generationConfig.temperature = request.temperature
        }
        if (typeof request.max_tokens === "number") {
            generationConfig.maxOutputTokens = request.max_tokens
        }

        const requestPayload: Record<string, unknown> = { contents }
        if (systemInstruction) {
            requestPayload.systemInstruction = { parts: [{ text: systemInstruction }] }
        }
        if (Object.keys(generationConfig).length > 0) {
            requestPayload.generationConfig = generationConfig
        }

        const body = JSON.stringify({
            project: projectId,
            model: request.model,
            request: requestPayload,
            userAgent: "antigravity",
            requestId: `agent-${randomUUID()}`
        })

        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}/v1internal:${GENERATE_ACTION}`, {
                    method: "POST",
                    headers: this.getAuthHeaders(authToken),
                    body,
                    signal: AbortSignal.timeout(120000)
                })

                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
                }

                const data = await response.json()
                return extractTextFromResponse(data)
            } 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")
    }

    /**
     * Get available Antigravity models
     */
    async getModels(accessToken?: string): Promise<AntigravityModel[]> {
        // Return full model list even if not authenticated, so they appear in selectors
        // Authenticaton is checked during actual chat requests
        const now = Date.now()
        if (this.modelsCache && (now - this.modelsCacheTime) < this.CACHE_TTL_MS) {
            return this.modelsCache
        }

        // If authenticated, return full model list
        this.modelsCache = ANTIGRAVITY_MODELS
        this.modelsCacheTime = now
        return ANTIGRAVITY_MODELS
    }

    /**
     * Test connection to Antigravity API
     */
    async testConnection(accessToken?: string, projectIdOverride?: string): Promise<{ connected: boolean; error?: string; status?: number }> {
        if (!this.isAuthenticated(accessToken)) {
            return { connected: false, error: "Not authenticated" }
        }

        try {
            const authToken = this.resolveAccessToken(accessToken)
            if (!authToken) {
                return { connected: false, error: "Not authenticated" }
            }

            const requestedProjectId = projectIdOverride?.trim()
            const loadEndpoints = Array.from(new Set([...LOAD_ENDPOINTS, ...this.config.endpoints]))
            let lastErrorText = ""
            let lastStatus: number | undefined

            const tryLoad = async (metadata: Record<string, string>): Promise<boolean> => {
                for (const endpoint of loadEndpoints) {
                    const response = await fetch(`${endpoint}/v1internal:loadCodeAssist`, {
                        method: "POST",
                        headers: this.getLoadHeaders(authToken),
                        body: JSON.stringify({ metadata }),
                        signal: AbortSignal.timeout(10000)
                    })
                    if (response.ok) {
                        return true
                    }
                    lastStatus = response.status
                    lastErrorText = await response.text().catch(() => "") || response.statusText
                }
                return false
            }

            const baseMetadata: Record<string, string> = { ...LOAD_ASSIST_METADATA }
            if (requestedProjectId) {
                baseMetadata.duetProject = requestedProjectId
            }
            let success = await tryLoad(baseMetadata)
            if (!success && !requestedProjectId) {
                const fallbackMetadata: Record<string, string> = {
                    ...LOAD_ASSIST_METADATA,
                    duetProject: DEFAULT_PROJECT_ID,
                }
                success = await tryLoad(fallbackMetadata)
            }
            if (success) {
                return { connected: true }
            }
            return {
                connected: false,
                status: lastStatus,
                error: lastErrorText || "Connection test failed"
            }
        } catch (error) {
            console.warn("Antigravity connection test failed:", error)
            return { connected: false, error: error instanceof Error ? error.message : String(error) }
        }
    }

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

        const content = await this.requestGenerateContent(request, accessToken, projectIdOverride)
        yield {
            id: randomUUID(),
            choices: [
                {
                    index: 0,
                    delta: { content },
                    finish_reason: "stop"
                }
            ]
        }
    }

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

        const content = await this.requestGenerateContent(request, accessToken, projectIdOverride)
        return {
            id: randomUUID(),
            choices: [
                {
                    index: 0,
                    message: {
                        role: "assistant",
                        content
                    },
                    finish_reason: "stop"
                }
            ]
        }
    }
}

export function getDefaultAntigravityConfig(): AntigravityConfig {
    return {
        enabled: true,
        endpoints: [
            "https://daily-cloudcode-pa.sandbox.googleapis.com",
            "https://autopush-cloudcode-pa.sandbox.googleapis.com",
            "https://cloudcode-pa.googleapis.com"
        ]
    }
}
