import { createHash, randomBytes, randomUUID } from "crypto"
import { createServer } from "http"
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 Authorization Code + PKCE configuration (Antigravity-compatible)
const GOOGLE_OAUTH_CONFIG = {
    clientId: process.env.ANTIGRAVITY_GOOGLE_CLIENT_ID || "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com",
    clientSecret: process.env.ANTIGRAVITY_GOOGLE_CLIENT_SECRET || "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",
    redirectUri: process.env.ANTIGRAVITY_GOOGLE_REDIRECT_URI || "http://localhost:51121/oauth-callback",
    authEndpoint: "https://accounts.google.com/o/oauth2/v2/auth",
    tokenEndpoint: "https://oauth2.googleapis.com/token",
    scopes: [
        "https://www.googleapis.com/auth/cloud-platform",
        "https://www.googleapis.com/auth/userinfo.email",
        "https://www.googleapis.com/auth/userinfo.profile",
        "https://www.googleapis.com/auth/cclog",
        "https://www.googleapis.com/auth/experimentsandconfigs",
    ],
}

const AUTH_SESSION_TTL_MS = 10 * 60 * 1000
const DEFAULT_POLL_INTERVAL_SEC = 5
const callbackUrl = new URL(GOOGLE_OAUTH_CONFIG.redirectUri)
const callbackPath = callbackUrl.pathname || "/oauth-callback"
const callbackPort = Number(callbackUrl.port || "0") || (callbackUrl.protocol === "https:" ? 443 : 80)

type OAuthSession = {
    verifier: string
    createdAt: number
    expiresAt: number
    token?: {
        accessToken: string
        refreshToken?: string
        expiresIn: number
        tokenType?: string
        scope?: string
    }
    error?: string
}

// Active OAuth sessions (in-memory, per-server instance)
const oauthSessions = new Map<string, OAuthSession>()
let oauthCallbackServer: ReturnType<typeof createServer> | null = null

function base64UrlEncode(value: Buffer): string {
    return value
        .toString("base64")
        .replace(/\+/g, "-")
        .replace(/\//g, "_")
        .replace(/=+$/g, "")
}

function createCodeVerifier(): string {
    return base64UrlEncode(randomBytes(32))
}

function createCodeChallenge(verifier: string): string {
    const digest = createHash("sha256").update(verifier).digest()
    return base64UrlEncode(digest)
}

function getAccessTokenFromHeader(authorization: string | undefined): string | null {
    if (!authorization) return null
    const [type, token] = authorization.split(" ")
    if (!type || type.toLowerCase() !== "bearer" || !token) return null
    return token.trim()
}

function getProjectIdFromHeader(value: string | string[] | undefined): string | undefined {
    if (typeof value === "string" && value.trim()) return value.trim()
    if (Array.isArray(value)) {
        const entry = value.find((item) => typeof item === "string" && item.trim())
        if (entry) return entry.trim()
    }
    return undefined
}

async function exchangeAuthorizationCode(code: string, verifier: string): Promise<{
    accessToken: string
    refreshToken?: string
    expiresIn: number
    tokenType?: string
    scope?: string
}> {
    const params = new URLSearchParams({
        client_id: GOOGLE_OAUTH_CONFIG.clientId,
        code,
        grant_type: "authorization_code",
        redirect_uri: GOOGLE_OAUTH_CONFIG.redirectUri,
        code_verifier: verifier,
    })
    if (GOOGLE_OAUTH_CONFIG.clientSecret) {
        params.set("client_secret", GOOGLE_OAUTH_CONFIG.clientSecret)
    }

    const response = await fetch(GOOGLE_OAUTH_CONFIG.tokenEndpoint, {
        method: "POST",
        headers: {
            "Content-Type": "application/x-www-form-urlencoded",
        },
        body: params,
    })

    if (!response.ok) {
        const errorText = await response.text()
        throw new Error(errorText || `Token exchange failed (${response.status})`)
    }

    const data = await response.json() as any
    return {
        accessToken: data.access_token,
        refreshToken: data.refresh_token,
        expiresIn: data.expires_in,
        tokenType: data.token_type,
        scope: data.scope,
    }
}

function ensureOAuthCallbackServer(logger: Logger): void {
    if (oauthCallbackServer) return
    oauthCallbackServer = createServer((req, res) => {
        void (async () => {
            try {
                const url = new URL(req.url || "", GOOGLE_OAUTH_CONFIG.redirectUri)
                if (url.pathname !== callbackPath) {
                    res.writeHead(404)
                    res.end()
                    return
                }

                const state = url.searchParams.get("state")
                const code = url.searchParams.get("code")
                const error = url.searchParams.get("error")

                if (!state) {
                    res.writeHead(400, { "Content-Type": "text/plain" })
                    res.end("Missing OAuth state.")
                    return
                }

                const session = oauthSessions.get(state)
                if (!session) {
                    res.writeHead(404, { "Content-Type": "text/plain" })
                    res.end("OAuth session not found or expired.")
                    return
                }

                if (error) {
                    session.error = error
                    res.writeHead(200, { "Content-Type": "text/html" })
                    res.end("<h2>Sign-in cancelled.</h2><p>You can close this window.</p>")
                    return
                }

                if (!code) {
                    session.error = "Missing authorization code."
                    res.writeHead(400, { "Content-Type": "text/plain" })
                    res.end("Missing authorization code.")
                    return
                }

                try {
                    const token = await exchangeAuthorizationCode(code, session.verifier)
                    session.token = token
                    session.error = undefined

                    res.writeHead(200, { "Content-Type": "text/html" })
                    res.end("<h2>Sign-in complete.</h2><p>You can close this window and return to the app.</p>")
                } catch (err) {
                    const message = err instanceof Error ? err.message : "OAuth callback failed."
                    session.error = message
                    res.writeHead(500, { "Content-Type": "text/plain" })
                    res.end(message)
                }
            } catch (err) {
                const message = err instanceof Error ? err.message : "OAuth callback failed."
                res.writeHead(500, { "Content-Type": "text/plain" })
                res.end(message)
            }
        })()
    })

    oauthCallbackServer.on("error", (err) => {
        logger.error({ err, port: callbackPort }, "Antigravity OAuth callback server failed to start")
        oauthCallbackServer = null
    })

    oauthCallbackServer.listen(callbackPort, "127.0.0.1", () => {
        logger.info({ port: callbackPort, path: callbackPath }, "Antigravity OAuth callback server listening")
    })
}

function cleanupExpiredSessions(): void {
    const now = Date.now()
    for (const [id, session] of oauthSessions) {
        if (session.expiresAt <= now) {
            oauthSessions.delete(id)
        }
    }
}

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 accessToken = getAccessTokenFromHeader(request.headers.authorization)
            const models = await client.getModels(accessToken ?? undefined)

            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 accessToken = getAccessTokenFromHeader(request.headers.authorization)
            const authenticated = client.isAuthenticated(accessToken ?? undefined)
            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 accessToken = getAccessTokenFromHeader(request.headers.authorization)
            const projectId = getProjectIdFromHeader(request.headers["x-antigravity-project"])
            const result = await client.testConnection(accessToken ?? undefined, projectId)
            return result
        } catch (error) {
            logger.error({ error }, "Antigravity connection test failed")
            return reply.status(500).send({ error: "Connection test failed" })
        }
    })

    // ==========================================
    // Google OAuth Authorization Flow (PKCE)
    // ==========================================

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

            const sessionId = randomUUID()
            const verifier = createCodeVerifier()
            const challenge = createCodeChallenge(verifier)

            const authUrl = new URL(GOOGLE_OAUTH_CONFIG.authEndpoint)
            authUrl.searchParams.set("client_id", GOOGLE_OAUTH_CONFIG.clientId)
            authUrl.searchParams.set("response_type", "code")
            authUrl.searchParams.set("redirect_uri", GOOGLE_OAUTH_CONFIG.redirectUri)
            authUrl.searchParams.set("scope", GOOGLE_OAUTH_CONFIG.scopes.join(" "))
            authUrl.searchParams.set("code_challenge", challenge)
            authUrl.searchParams.set("code_challenge_method", "S256")
            authUrl.searchParams.set("state", sessionId)
            authUrl.searchParams.set("access_type", "offline")
            authUrl.searchParams.set("prompt", "consent")

            oauthSessions.set(sessionId, {
                verifier,
                createdAt: Date.now(),
                expiresAt: Date.now() + AUTH_SESSION_TTL_MS,
            })
            cleanupExpiredSessions()

            return {
                sessionId,
                userCode: "",
                verificationUrl: authUrl.toString(),
                expiresIn: Math.floor(AUTH_SESSION_TTL_MS / 1000),
                interval: DEFAULT_POLL_INTERVAL_SEC,
            }
        } catch (error: any) {
            logger.error({ error: error.message, stack: error.stack }, "Failed to start OAuth authorization")
            return reply.status(500).send({
                error: "Failed to start authentication",
                details: error.message,
            })
        }
    })

    // Step 2: Poll for token (called by client after browser sign-in)
    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" })
            }

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

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

            if (session.error) {
                oauthSessions.delete(sessionId)
                return { status: "error", error: session.error }
            }

            if (!session.token) {
                return { status: "pending", interval: DEFAULT_POLL_INTERVAL_SEC }
            }

            const token = session.token
            oauthSessions.delete(sessionId)

            return {
                status: "success",
                accessToken: token.accessToken,
                refreshToken: token.refreshToken,
                expiresIn: token.expiresIn,
                tokenType: token.tokenType,
                scope: token.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 params = new URLSearchParams({
                client_id: GOOGLE_OAUTH_CONFIG.clientId,
                refresh_token: refreshToken,
                grant_type: "refresh_token",
            })
            if (GOOGLE_OAUTH_CONFIG.clientSecret) {
                params.set("client_secret", GOOGLE_OAUTH_CONFIG.clientSecret)
            }

            const response = await fetch(GOOGLE_OAUTH_CONFIG.tokenEndpoint, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/x-www-form-urlencoded'
                },
                body: params
            })

            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
            }
            const accessToken = getAccessTokenFromHeader(request.headers.authorization)
            const projectId = getProjectIdFromHeader(request.headers["x-antigravity-project"])

            // 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,
                        accessToken ?? undefined,
                        projectId,
                        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,
                    accessToken ?? undefined,
                    projectId,
                    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 OAuth flow!")
}

/**
 * Streaming chat with tool execution loop
 */
async function streamWithToolLoop(
    client: AntigravityClient,
    request: ChatRequest,
    accessToken: string | undefined,
    projectId: string | undefined,
    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 }, accessToken, projectId)) {
            // 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,
    accessToken: string | undefined,
    projectId: string | undefined,
    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 }, accessToken, projectId)
        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
}
