/*
 * Copyright 2025 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.hbasesoft.framework.ai.agent.tool.browser;

import java.util.Arrays;
import java.util.List;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Primary;
import org.springframework.stereotype.Service;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.hbasesoft.framework.ai.agent.config.IManusProperties;
import com.hbasesoft.framework.ai.agent.config.ManusProperties;
import com.hbasesoft.framework.ai.agent.tool.filesystem.UnifiedDirectoryManager;
import com.hbasesoft.framework.ai.agent.tool.innerStorage.SmartContentSavingService;
import com.microsoft.playwright.Browser;
import com.microsoft.playwright.BrowserType;
import com.microsoft.playwright.Page;
import com.microsoft.playwright.Playwright;

import jakarta.annotation.PreDestroy;

@Service
@Primary
public class ChromeDriverService implements IChromeDriverService {

    private static final Logger log = LoggerFactory.getLogger(ChromeDriverService.class);

    private final ConcurrentHashMap<String, DriverWrapper> drivers = new ConcurrentHashMap<>();

    private final Lock driverLock = new ReentrantLock();

    private IManusProperties manusProperties;

    private SmartContentSavingService innerStorageService;

    private UnifiedDirectoryManager unifiedDirectoryManager;

    private final ObjectMapper objectMapper;

    @Autowired(required = false)
    private SpringBootPlaywrightInitializer playwrightInitializer;

    /**
     * Shared directory for storing cookies
     */
    /**
     * Shared directory for storing cookies
     */
    private String sharedDir;

    /**
     * Get current shared directory
     */
    public String getSharedDir() {
        return sharedDir;
    }

    /**
     * Save all cookies from drivers to global shared directory (cookies.json)
     */
    public void saveCookiesToSharedDir() {
        // Get the first available driver
        DriverWrapper driver = drivers.values().stream().findFirst().orElse(null);
        if (driver == null) {
            log.warn("No driver found for saving cookies");
            return;
        }
        try {
            List<com.microsoft.playwright.options.Cookie> cookies = driver.getCurrentPage().context().cookies();
            String cookieFile = sharedDir + "/cookies.json";
            try (java.io.FileWriter writer = new java.io.FileWriter(cookieFile)) {
                writer.write(objectMapper.writeValueAsString(cookies));
            }
            log.info("Cookies saved to {}", cookieFile);
        }
        catch (Exception e) {
            log.error("Failed to save cookies", e);
        }
    }

    /**
     * Load cookies from global shared directory to all drivers
     */
    public void loadCookiesFromSharedDir() {
        String cookieFile = sharedDir + "/cookies.json";
        java.io.File file = new java.io.File(cookieFile);
        if (!file.exists()) {
            log.warn("Cookie file does not exist: {}", cookieFile);
            return;
        }
        try (java.io.FileReader reader = new java.io.FileReader(cookieFile)) {
            // Replace FastJSON's JSON.parseArray with Jackson's objectMapper.readValue
            List<com.microsoft.playwright.options.Cookie> cookies = objectMapper.readValue(reader,
                new TypeReference<List<com.microsoft.playwright.options.Cookie>>() {
                });
            for (DriverWrapper driver : drivers.values()) {
                driver.getCurrentPage().context().addCookies(cookies);
            }
            log.info("Cookies loaded from {} to all drivers", cookieFile);
        }
        catch (Exception e) {
            log.error("Failed to load cookies for all drivers", e);
        }
    }

    public ChromeDriverService(IManusProperties manusProperties, SmartContentSavingService innerStorageService,
        UnifiedDirectoryManager unifiedDirectoryManager, ObjectMapper objectMapper) {
        this.manusProperties = manusProperties;
        this.innerStorageService = innerStorageService;
        this.unifiedDirectoryManager = unifiedDirectoryManager;
        this.objectMapper = objectMapper;
    }

    public void init() {
        // Use UnifiedDirectoryManager to get the shared directory for playwright
        try {
            java.nio.file.Path playwrightDir = unifiedDirectoryManager.getWorkingDirectory().resolve("playwright");
            unifiedDirectoryManager.ensureDirectoryExists(playwrightDir);
            this.sharedDir = playwrightDir.toString();
        }
        catch (java.io.IOException e) {
            log.error("Failed to create playwright directory", e);
            this.sharedDir = unifiedDirectoryManager.getWorkingDirectory().resolve("playwright").toString();
        }
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            log.info("JVM shutting down - cleaning up Playwright processes");
            cleanupAllPlaywrightProcesses();
        }));
    }

    public DriverWrapper getDriver(String planId) {
        if (planId == null) {
            throw new IllegalArgumentException("planId cannot be null");
        }

        DriverWrapper currentDriver = drivers.get(planId);
        if (currentDriver != null) {
            return currentDriver;
        }

        try {
            driverLock.lock();
            currentDriver = drivers.get(planId);
            if (currentDriver != null) {
                return currentDriver;
            }
            log.info("Creating new Playwright Browser instance for planId: {}", planId);
            currentDriver = createNewDriver(); // createNewDriver will now pass sharedDir
            if (currentDriver != null) { // Check if driver creation was successful
                drivers.put(planId, currentDriver);
            }
            else {
                // Handle the case where driver creation failed, e.g., log an error or
                // throw an exception
                log.error("Failed to create new driver for planId: {}. createNewDriver returned null.", planId);
                // Optionally throw an exception to indicate failure to the caller
                // throw new RuntimeException("Failed to create new driver for planId: " +
                // planId);
            }
        }
        finally {
            driverLock.unlock();
        }

        return currentDriver;
    }

    private void cleanupAllPlaywrightProcesses() {
        try {
            drivers.clear();
            log.info("Successfully cleaned up all Playwright processes	");
        }
        catch (Exception e) {
            log.error("Error cleaning up Browser processes", e);
        }
    }

    public void closeDriverForPlan(String planId) {
        DriverWrapper driver = drivers.remove(planId);
        if (driver != null) {
            driver.close();
        }
    }

    private DriverWrapper createNewDriver() {
        log.info("Creating new browser driver");
        return createDriverInstance();
    }

    /**
     * Create browser driver instance
     */
    private DriverWrapper createDriverInstance() {
        // Set system properties for Playwright configuration
        System.setProperty("playwright.browsers.path", System.getProperty("user.home") + "/.cache/ms-playwright");

        // Set custom driver temp directory to avoid classpath issues
        System.setProperty("playwright.driver.tmpdir", System.getProperty("java.io.tmpdir"));

        // Skip browser download if browsers are already installed
        System.setProperty("PLAYWRIGHT_SKIP_BROWSER_DOWNLOAD", "1");

        // Try to create Playwright instance using Spring Boot initializer
        Playwright playwright = null;
        try {
            if (playwrightInitializer != null && playwrightInitializer.canInitialize()) {
                log.info("Using SpringBootPlaywrightInitializer");
                playwright = playwrightInitializer.createPlaywright();
            }
            else {
                log.info("Using standard Playwright initialization");
                playwright = Playwright.create();
            }

            // Get browser type
            BrowserType browserType = getBrowserTypeFromEnv(playwright);
            log.info("Using browser type: {}", browserType.name());

            BrowserType.LaunchOptions options = new BrowserType.LaunchOptions();

            // Basic configuration
            options.setArgs(Arrays.asList("--remote-allow-origins=*", "--disable-blink-features=AutomationControlled",
                "--disable-infobars", "--disable-notifications", "--disable-dev-shm-usage", "--lang=zh-CN,zh,en-US,en",
                "--user-agent=" + getRandomUserAgent(), "--window-size=1920,1080"));

            // Decide whether to use headless mode based on configuration
            if (manusProperties.getBrowserHeadless()) {
                log.info("Enable Playwright headless mode");
                options.setHeadless(true);
            }
            else {
                log.info("Enable Playwright non-headless mode");
                options.setHeadless(false);
            }

            Browser browser = browserType.launch(options);
            log.info("Created new Playwright Browser instance");

            // Create new page and configure timeout
            Page page = browser.newPage();

            // Set default timeout based on configuration
            Integer timeout = manusProperties.getBrowserRequestTimeout();
            if (timeout != null && timeout > 0) {
                log.info("Setting browser page timeout to {} seconds", timeout);
                page.setDefaultTimeout(timeout * 1000); // Convert to milliseconds
            }

            return new DriverWrapper(playwright, browser, page, this.sharedDir, objectMapper);
        }
        catch (Exception e) {
            if (playwright != null) {
                try {
                    playwright.close();
                }
                catch (Exception ex) {
                    log.warn("Failed to close failed Playwright instance", ex);
                }
            }
            throw new RuntimeException("Failed to initialize Playwright Browser", e);
        }
    }

    /**
     * Get browser type, supports environment variable configuration
     */
    private BrowserType getBrowserTypeFromEnv(Playwright playwright) {
        String browserName = System.getenv("BROWSER");
        if (browserName == null) {
            browserName = "chromium";
        }

        switch (browserName.toLowerCase()) {
            case "webkit":
                return playwright.webkit();
            case "firefox":
                return playwright.firefox();
            case "chromium":
            default:
                return playwright.chromium();
        }
    }

    private String getRandomUserAgent() {
        List<String> userAgents = Arrays.asList(
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
            "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36 Edg/119.0.0.0");
        return userAgents.get(new Random().nextInt(userAgents.size()));
    }

    @PreDestroy
    public void cleanup() {
        log.info("Spring container shutting down - cleaning up Browser resources");
        cleanupAllPlaywrightProcesses();
    }

    public void setManusProperties(IManusProperties manusProperties) {
        this.manusProperties = (ManusProperties) manusProperties;
    }

    public IManusProperties getManusProperties() {
        return manusProperties;
    }

    public SmartContentSavingService getInnerStorageService() {
        return innerStorageService;
    }

    public UnifiedDirectoryManager getUnifiedDirectoryManager() {
        return unifiedDirectoryManager;
    }

}
