import { FastifyInstance } from "fastify"
import { AntigravityClient, type ChatRequest, getDefaultAntigravityConfig, type ChatMessage } from "../../integrations/antigravity"
import { Logger } from "../../logger"
import { CORE_TOOLS, executeTools, type ToolCall, type ToolResult } from "../../tools/executor"
import { getMcpManager } from "../../mcp/client"

interface AntigravityRouteDeps {
    logger: Logger
}

// Maximum number of tool execution loops
const MAX_TOOL_LOOPS = 10

// Google OAuth Device Flow configuration
// Using the same client ID as gcloud CLI / Cloud SDK
const GOOGLE_OAUTH_CONFIG = {
    clientId: "764086051850-6qr4p6gpi6hn506pt8ejuq83di341hur.apps.googleusercontent.com",
    clientSecret: "d-FL95Q19q7MQmFpd7hHD0Ty", // Public client secret for device flow
    deviceAuthEndpoint: "https://oauth2.googleapis.com/device/code",
    tokenEndpoint: "https://oauth2.googleapis.com/token",
    scopes: [
        "openid",
        "email",
        "profile",
        "https://www.googleapis.com/auth/cloud-platform"
    ]
}

// Active device auth sessions (in-memory, per-server instance)
const deviceAuthSessions = new Map<string, {
    deviceCode: string
    userCode: string
    verificationUrl: string
    expiresAt: number
    interval: number
}>()

export async function registerAntigravityRoutes(
    app: FastifyInstance,
    deps: AntigravityRouteDeps
) {
    const logger = deps.logger.child({ component: "antigravity-routes" })

    // Create shared client
    const client = new AntigravityClient(getDefaultAntigravityConfig())

    // List available Antigravity models
    app.get('/api/antigravity/models', async (request, reply) => {
        try {
            const models = await client.getModels()

            return {
                models: models.map(m => ({
                    id: m.id,
                    name: m.name,
                    family: m.family,
                    provider: "antigravity",
                    free: false, // These require Google OAuth
                    reasoning: m.reasoning,
                    tool_call: m.tool_call,
                    limit: m.limit
                }))
            }
        } catch (error) {
            logger.error({ error }, "Failed to list Antigravity models")
            return reply.status(500).send({ error: "Failed to list models" })
        }
    })

    // Check authentication status
    app.get('/api/antigravity/auth-status', async (request, reply) => {
        try {
            const authenticated = client.isAuthenticated()
            return { authenticated }
        } catch (error) {
            logger.error({ error }, "Antigravity auth status check failed")
            return reply.status(500).send({ error: "Auth status check failed" })
        }
    })

    // Test connection
    app.get('/api/antigravity/test', async (request, reply) => {
        try {
            const connected = await client.testConnection()
            return { connected }
        } catch (error) {
            logger.error({ error }, "Antigravity connection test failed")
            return reply.status(500).send({ error: "Connection test failed" })
        }
    })

    // ==========================================
    // Google Device Authorization Flow Endpoints
    // ==========================================

    // Step 1: Start device authorization - returns user_code and verification URL
    app.post('/api/antigravity/device-auth/start', async (request, reply) => {
        try {
            logger.info("Starting Google Device Authorization flow for Antigravity")

            const response = await fetch(GOOGLE_OAUTH_CONFIG.deviceAuthEndpoint, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/x-www-form-urlencoded'
                },
                body: new URLSearchParams({
                    client_id: GOOGLE_OAUTH_CONFIG.clientId,
                    scope: GOOGLE_OAUTH_CONFIG.scopes.join(' ')
                })
            })

            if (!response.ok) {
                const errorText = await response.text()
                logger.error({ error: errorText, status: response.status }, "Device auth request failed")
                return reply.status(500).send({
                    error: `Device authorization failed: ${response.status}`,
                    details: errorText
                })
            }

            const data = await response.json() as {
                device_code: string
                user_code: string
                verification_url: string
                expires_in: number
                interval: number
            }

            // Generate a session ID for tracking this auth flow
            const sessionId = crypto.randomUUID()

            // Store the session
            deviceAuthSessions.set(sessionId, {
                deviceCode: data.device_code,
                userCode: data.user_code,
                verificationUrl: data.verification_url,
                expiresAt: Date.now() + (data.expires_in * 1000),
                interval: data.interval
            })

            // Clean up expired sessions
            for (const [id, session] of deviceAuthSessions) {
                if (session.expiresAt < Date.now()) {
                    deviceAuthSessions.delete(id)
                }
            }

            logger.info({ sessionId, userCode: data.user_code, verificationUrl: data.verification_url }, "Device auth session created")

            return {
                sessionId,
                userCode: data.user_code,
                verificationUrl: data.verification_url,
                expiresIn: data.expires_in,
                interval: data.interval
            }
        } catch (error: any) {
            logger.error({ error: error.message, stack: error.stack }, "Failed to start device authorization")
            return reply.status(500).send({
                error: "Failed to start device authorization",
                details: error.message
            })
        }
    })

    // Step 2: Poll for token (called by client after user enters code)
    app.post('/api/antigravity/device-auth/poll', async (request, reply) => {
        try {
            const { sessionId } = request.body as { sessionId: string }

            if (!sessionId) {
                return reply.status(400).send({ error: "Missing sessionId" })
            }

            const session = deviceAuthSessions.get(sessionId)
            if (!session) {
                return reply.status(404).send({ error: "Session not found or expired" })
            }

            if (session.expiresAt < Date.now()) {
                deviceAuthSessions.delete(sessionId)
                return reply.status(410).send({ error: "Session expired" })
            }

            // Poll Google's token endpoint
            const response = await fetch(GOOGLE_OAUTH_CONFIG.tokenEndpoint, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/x-www-form-urlencoded'
                },
                body: new URLSearchParams({
                    client_id: GOOGLE_OAUTH_CONFIG.clientId,
                    client_secret: GOOGLE_OAUTH_CONFIG.clientSecret,
                    device_code: session.deviceCode,
                    grant_type: 'urn:ietf:params:oauth:grant-type:device_code'
                })
            })

            const data = await response.json() as any

            if (data.error) {
                // Still waiting for user
                if (data.error === 'authorization_pending') {
                    return { status: 'pending', interval: session.interval }
                }
                // Slow down polling
                if (data.error === 'slow_down') {
                    session.interval = Math.min(session.interval + 5, 60)
                    return { status: 'pending', interval: session.interval }
                }
                // User denied or other error
                if (data.error === 'access_denied') {
                    deviceAuthSessions.delete(sessionId)
                    return { status: 'denied' }
                }
                if (data.error === 'expired_token') {
                    deviceAuthSessions.delete(sessionId)
                    return { status: 'expired' }
                }

                logger.error({ error: data.error }, "Token poll error")
                return { status: 'error', error: data.error }
            }

            // Success! We have tokens
            deviceAuthSessions.delete(sessionId)

            logger.info("Device authorization successful")

            return {
                status: 'success',
                accessToken: data.access_token,
                refreshToken: data.refresh_token,
                expiresIn: data.expires_in,
                tokenType: data.token_type,
                scope: data.scope
            }
        } catch (error) {
            logger.error({ error }, "Failed to poll for token")
            return reply.status(500).send({ error: "Failed to poll for token" })
        }
    })

    // Refresh an expired token
    app.post('/api/antigravity/device-auth/refresh', async (request, reply) => {
        try {
            const { refreshToken } = request.body as { refreshToken: string }

            if (!refreshToken) {
                return reply.status(400).send({ error: "Missing refreshToken" })
            }

            const response = await fetch(GOOGLE_OAUTH_CONFIG.tokenEndpoint, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/x-www-form-urlencoded'
                },
                body: new URLSearchParams({
                    client_id: GOOGLE_OAUTH_CONFIG.clientId,
                    client_secret: GOOGLE_OAUTH_CONFIG.clientSecret,
                    refresh_token: refreshToken,
                    grant_type: 'refresh_token'
                })
            })

            if (!response.ok) {
                const error = await response.text()
                logger.error({ error }, "Token refresh failed")
                return reply.status(401).send({ error: "Token refresh failed" })
            }

            const data = await response.json() as any

            return {
                accessToken: data.access_token,
                expiresIn: data.expires_in,
                tokenType: data.token_type
            }
        } catch (error) {
            logger.error({ error }, "Failed to refresh token")
            return reply.status(500).send({ error: "Failed to refresh token" })
        }
    })

    // Chat completion endpoint WITH MCP TOOL SUPPORT
    app.post('/api/antigravity/chat', async (request, reply) => {
        try {
            const chatRequest = request.body as ChatRequest & {
                workspacePath?: string
                enableTools?: boolean
            }

            // Extract workspace path for tool execution
            const workspacePath = chatRequest.workspacePath || process.cwd()
            const enableTools = chatRequest.enableTools !== false

            logger.info({
                workspacePath,
                receivedWorkspacePath: chatRequest.workspacePath,
                enableTools,
                model: chatRequest.model
            }, "Antigravity chat request received")

            // Handle streaming with tool loop
            if (chatRequest.stream) {
                reply.raw.writeHead(200, {
                    'Content-Type': 'text/event-stream',
                    'Cache-Control': 'no-cache',
                    'Connection': 'keep-alive',
                })

                try {
                    await streamWithToolLoop(
                        client,
                        chatRequest,
                        workspacePath,
                        enableTools,
                        reply.raw,
                        logger
                    )
                    reply.raw.end()
                } catch (streamError) {
                    logger.error({ error: streamError }, "Antigravity streaming failed")
                    reply.raw.write(`data: ${JSON.stringify({ error: String(streamError) })}\n\n`)
                    reply.raw.end()
                }
            } else {
                // Non-streaming with tool loop
                const response = await chatWithToolLoop(
                    client,
                    chatRequest,
                    workspacePath,
                    enableTools,
                    logger
                )
                return response
            }
        } catch (error) {
            logger.error({ error }, "Antigravity chat request failed")
            return reply.status(500).send({ error: "Chat request failed" })
        }
    })

    logger.info("Antigravity routes registered with Google Device Auth flow!")
}

/**
 * Streaming chat with tool execution loop
 */
async function streamWithToolLoop(
    client: AntigravityClient,
    request: ChatRequest,
    workspacePath: string,
    enableTools: boolean,
    rawResponse: any,
    logger: Logger
): Promise<void> {
    let messages = [...request.messages]
    let loopCount = 0

    // Load MCP tools from workspace config
    let allTools = [...CORE_TOOLS]
    if (enableTools && workspacePath) {
        try {
            const mcpManager = getMcpManager()
            await mcpManager.loadConfig(workspacePath)
            const mcpTools = await mcpManager.getToolsAsOpenAIFormat()
            allTools = [...CORE_TOOLS, ...mcpTools]
            if (mcpTools.length > 0) {
                logger.info({ mcpToolCount: mcpTools.length }, "Loaded MCP tools for Antigravity")
            }
        } catch (mcpError) {
            logger.warn({ error: mcpError }, "Failed to load MCP tools")
        }
    }

    // Inject tools if enabled
    const requestWithTools: ChatRequest = {
        ...request,
        tools: enableTools ? allTools : undefined,
        tool_choice: enableTools ? "auto" : undefined
    }

    while (loopCount < MAX_TOOL_LOOPS) {
        loopCount++

        // Accumulate tool calls from stream
        let accumulatedToolCalls: { [index: number]: { id: string; name: string; arguments: string } } = {}
        let hasToolCalls = false
        let textContent = ""

        // Stream response
        for await (const chunk of client.chatStream({ ...requestWithTools, messages })) {
            // Write chunk to client
            rawResponse.write(`data: ${JSON.stringify(chunk)}\n\n`)

            const choice = chunk.choices[0]
            if (!choice) continue

            // Accumulate text content
            if (choice.delta?.content) {
                textContent += choice.delta.content
            }

            // Accumulate tool calls from delta (if API supports it)
            const deltaToolCalls = (choice.delta as any)?.tool_calls
            if (deltaToolCalls) {
                hasToolCalls = true
                for (const tc of deltaToolCalls) {
                    const idx = tc.index ?? 0
                    if (!accumulatedToolCalls[idx]) {
                        accumulatedToolCalls[idx] = { id: tc.id || "", name: "", arguments: "" }
                    }
                    if (tc.id) accumulatedToolCalls[idx].id = tc.id
                    if (tc.function?.name) accumulatedToolCalls[idx].name += tc.function.name
                    if (tc.function?.arguments) accumulatedToolCalls[idx].arguments += tc.function.arguments
                }
            }

            // Check if we should stop
            if (choice.finish_reason === "stop") {
                rawResponse.write('data: [DONE]\n\n')
                return
            }
        }

        // If no tool calls, we're done
        if (!hasToolCalls || !enableTools) {
            rawResponse.write('data: [DONE]\n\n')
            return
        }

        // Convert accumulated tool calls
        const toolCalls: ToolCall[] = Object.values(accumulatedToolCalls).map(tc => ({
            id: tc.id,
            type: "function" as const,
            function: {
                name: tc.name,
                arguments: tc.arguments
            }
        }))

        if (toolCalls.length === 0) {
            rawResponse.write('data: [DONE]\n\n')
            return
        }

        logger.info({ toolCalls: toolCalls.map(tc => tc.function.name) }, "Executing tool calls")

        // Add assistant message with tool calls
        const assistantMessage: ChatMessage = {
            role: "assistant",
            content: textContent || undefined,
            tool_calls: toolCalls
        }
        messages.push(assistantMessage)

        // Execute tools
        const toolResults = await executeTools(workspacePath, toolCalls)

        // Notify client about tool execution via special event
        for (const result of toolResults) {
            const toolEvent = {
                type: "tool_result",
                tool_call_id: result.tool_call_id,
                content: result.content
            }
            rawResponse.write(`data: ${JSON.stringify(toolEvent)}\n\n`)
        }

        // Add tool results to messages
        for (const result of toolResults) {
            const toolMessage: ChatMessage = {
                role: "tool",
                content: result.content,
                tool_call_id: result.tool_call_id
            }
            messages.push(toolMessage)
        }

        logger.info({ loopCount, toolsExecuted: toolResults.length }, "Tool loop iteration complete")
    }

    logger.warn({ loopCount }, "Max tool loops reached")
    rawResponse.write('data: [DONE]\n\n')
}

/**
 * Non-streaming chat with tool execution loop
 */
async function chatWithToolLoop(
    client: AntigravityClient,
    request: ChatRequest,
    workspacePath: string,
    enableTools: boolean,
    logger: Logger
): Promise<any> {
    let messages = [...request.messages]
    let loopCount = 0
    let lastResponse: any = null

    // Inject tools if enabled
    const requestWithTools: ChatRequest = {
        ...request,
        tools: enableTools ? CORE_TOOLS : undefined,
        tool_choice: enableTools ? "auto" : undefined
    }

    while (loopCount < MAX_TOOL_LOOPS) {
        loopCount++

        const response = await client.chat({ ...requestWithTools, messages, stream: false })
        lastResponse = response

        const choice = response.choices[0]
        if (!choice) break

        const toolCalls = (choice.message as any)?.tool_calls

        // If no tool calls, return
        if (!toolCalls || toolCalls.length === 0 || !enableTools) {
            return response
        }

        logger.info({ toolCalls: toolCalls.map((tc: any) => tc.function.name) }, "Executing tool calls")

        // Add assistant message
        const assistantMessage: ChatMessage = {
            role: "assistant",
            content: (choice.message as any).content || undefined,
            tool_calls: toolCalls
        }
        messages.push(assistantMessage)

        // Execute tools
        const toolResults = await executeTools(workspacePath, toolCalls)

        // Add tool results
        for (const result of toolResults) {
            const toolMessage: ChatMessage = {
                role: "tool",
                content: result.content,
                tool_call_id: result.tool_call_id
            }
            messages.push(toolMessage)
        }

        logger.info({ loopCount, toolsExecuted: toolResults.length }, "Tool loop iteration complete")
    }

    logger.warn({ loopCount }, "Max tool loops reached")
    return lastResponse
}
