from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from bs4 import BeautifulSoup
import re
import logging
import json
import time

# Set up basic logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

DOMAIN = "https://www.globalfirepower.com"


class GlobalFirePowerScraper:
    def __init__(self, headless=True, debug=False):
        self.debug = debug
        self.headless = headless
        self._setup_driver()

    def _setup_driver(self):
        """Set up the Selenium WebDriver"""
        chrome_options = Options()
        if self.headless:
            chrome_options.add_argument("--headless")
        chrome_options.add_argument("--no-sandbox")
        chrome_options.add_argument("--disable-dev-shm-usage")
        chrome_options.add_argument("--disable-gpu")
        chrome_options.add_argument("--window-size=1920,1080")
        chrome_options.add_argument("--start-maximized")

        try:
            from webdriver_manager.chrome import ChromeDriverManager
            service = Service(ChromeDriverManager().install())
            self.driver = webdriver.Chrome(service=service, options=chrome_options)
        except ImportError:
            self.driver = webdriver.Chrome(options=chrome_options)

        self.driver.implicitly_wait(10)

    def close(self):
        """Close the WebDriver"""
        if hasattr(self, 'driver'):
            self.driver.quit()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def search(self, country_id: str, save_to_file: bool = False) -> dict:
        """Search for country details by ID"""
        html = self._get_page(country_id)
        if html:
            results = self._extract_content(html)
            # Save results to a file if requested
            if save_to_file:
                self.save_to_json(results, f"{country_id}_military_data.json")
            return results
        return {}

    def get_raw_html(self, country_id: str) -> str | None:
        """Get the raw HTML content of a country's detail page for debugging"""
        return self._get_page(country_id)

    def get_raw_data(self, country_id: str) -> dict:
        """Get the raw data in a structured format without saving to a file"""
        html = self._get_page(country_id)
        if html:
            return self._extract_content(html)
        return {}

    def _get_page(self, country_id) -> str | None:
        """Get the HTML content of a country's detail page using Selenium"""
        url = f"{DOMAIN}/country-military-strength-detail.php?country_id={country_id}"
        try:
            self.driver.get(url)

            # Wait for the page to load completely by waiting for the collapsible section buttons
            WebDriverWait(self.driver, 20).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, "button.collapsible"))
            )

            # Scroll to ensure all elements are loaded
            self.driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
            time.sleep(1)
            self.driver.execute_script("window.scrollTo(0, 0);")
            time.sleep(1)

            # Find all collapsible buttons and expand them with better waiting
            buttons = WebDriverWait(self.driver, 15).until(
                EC.presence_of_all_elements_located((By.CSS_SELECTOR, "button.collapsible"))
            )

            for button in buttons:
                try:
                    # Scroll to the button with offset to account for headers
                    self.driver.execute_script("arguments[0].scrollIntoView({block: 'center', behavior: 'smooth'});",
                                               button)
                    time.sleep(0.5)

                    # Wait for button to be clickable
                    WebDriverWait(self.driver, 5).until(
                        EC.element_to_be_clickable(button)
                    )

                    # Click to expand
                    button.click()
                    logger.info(f"Clicked section: {button.text}")

                    # Wait for content to load
                    time.sleep(0.5)

                except Exception as e:
                    logger.warning(f"Failed to expand section {button.text}: {e}")
                    continue

            # Additional wait after expanding all sections
            time.sleep(2)

            # Get the page source after expanding all sections
            html = self.driver.page_source

            if self.debug:
                with open(f"debug_{country_id}.html", "w", encoding="utf-8") as f:
                    f.write(html)

            return html
        except Exception as e:
            logger.error(f"Error fetching page: {e}")
            return None

    def _extract_content(self, html) -> dict:
        """Extract search results from HTML"""
        soup = BeautifulSoup(html, 'html.parser')
        results = {}

        # Find all section headers
        buttons = soup.find_all("button", class_="collapsible")

        if not buttons:
            logger.warning("No collapsible buttons found in the page")
            return results

        for button in buttons:
            header_text = button.get_text(strip=True).rstrip(' [+]').lower()
            logger.info(f"Processing section: {header_text}")

            # Find the content div that follows this button
            content_div = button.find_next_sibling("div", class_="contentSpecs")

            if not content_div:
                logger.warning(f"No content div found for section: {header_text}")
                continue

            try:
                self._process_section(header_text, content_div, results)
            except Exception as e:
                logger.error(f"Error processing section '{header_text}': {e}")

        # Get comparable and neighboring powers
        try:
            comparable = soup.select(".moreLikePanel .textLargest")
            results['comparable_powers'] = [country.get_text(strip=True) for country in comparable]

            neighboring = soup.select(".neighborPanel .textLargest")
            results['neighbouring_powers'] = [country.get_text(strip=True) for country in neighboring]
        except Exception as e:
            logger.error(f"Error extracting powers lists: {e}")

        return results

    def _process_section(self, header, content_div, results):
        """Process each section based on its header"""
        # Normalize header text for matching
        normalized_header = header.lower().replace("-", " ").strip()

        if "at a glance" in normalized_header:
            results['at_a_glance'] = self._extract_at_a_glance(content_div)
        elif normalized_header == "overview":
            results['overview'] = self._extract_overview(content_div)
        elif normalized_header == "financials":
            results['financials'] = self._extract_statistics_section(
                content_div,
                description="Financial metrics related to defense spending and economic indicators.",
                metric="All monetary values presented in United States Dollar (USD$)."
            )
        elif normalized_header == "geography":
            results['geography'] = self._extract_statistics_section(
                content_div,
                description="Geographic metrics including land area, coastline, and shared borders.",
                metric="All distances presented in kilometers (km)."
            )
        elif normalized_header == "manpower":
            results['manpower'] = self._extract_statistics_section(
                content_div,
                description="The following values detail the maximum, theoretical persons the nation can commit to a war effort."
            )
        elif normalized_header == "airpower":
            results['airpower'] = self._extract_statistics_section(
                content_div,
                description="Aerial warfare capabilities including aircraft inventory and readiness."
            )
        elif normalized_header in ["land forces", "landforces"]:
            results['land_forces'] = self._extract_statistics_section(
                content_div,
                description="Land warfare capabilities including tanks, armored vehicles, and artillery."
            )
        elif normalized_header in ["naval forces", "navalforces"]:
            naval_stats = self._extract_statistics_section(
                content_div,
                description="Naval warfare capabilities including ships and submarines."
            )
            # Extract hull class descriptions if available
            hull_descriptions = {}
            hull_containers = content_div.select(".hullClassContainers")
            for hull in hull_containers:
                hull_type = hull.select_one(".textNormal")
                hull_desc = hull.select_one(".textSmall1")
                if hull_type and hull_desc:
                    hull_descriptions[hull_type.get_text(strip=True)] = hull_desc.get_text(strip=True)
            if hull_descriptions:
                naval_stats['hull_types'] = hull_descriptions
            results['naval_forces'] = naval_stats
        elif "end use" in normalized_header or "end-use" in normalized_header:
            products = [p.get_text(strip=True) for p in content_div.select('.prodTitleContainer')]
            results['end_use_products'] = {
                'description': "Industries that would become stressed in the event of Total War.",
                'products': products
            }
        elif "natural resources" in normalized_header:
            results['natural_resources'] = self._extract_statistics_section(
                content_div,
                description="Natural resources available to the country."
            )
        elif normalized_header == "logistics":
            results['logistics'] = self._extract_statistics_section(
                content_div,
                description="Logistics capabilities including transportation infrastructure."
            )
        else:
            logger.info(f"Skipping unhandled section: {header}")

    def _extract_at_a_glance(self, content_div):
        """Extract at-a-glance information"""
        result = {}

        # Extract graph data
        result['graph_data'] = self._get_graph_data_from_page()

        # Extract country info
        country_info = content_div.select_one(".textNormal")
        if country_info:
            result['country_info'] = country_info.get_text(strip=True)

        # Extract quick facts
        quick_facts = []
        for fact in content_div.select('.glanceDescription'):
            quick_facts.append(fact.get_text(strip=True))
        if quick_facts:
            result['quick_facts'] = quick_facts

        return result

    def _extract_overview(self, content_div):
        """Extract overview information"""
        result = {}

        # Extract rankings
        ranks = {}
        for ranking in content_div.select('.rankBaseContainer'):
            title = ranking.select_one(".textNormal")
            value = ranking.select_one(".textJumbo")
            suffix = ranking.select_one(".textSmall1")

            if title and value:
                rank_text = value.get_text(strip=True)
                if suffix:
                    rank_text += " " + suffix.get_text(strip=True)
                ranks[title.get_text(strip=True)] = rank_text

        if ranks:
            result['ranks'] = ranks

        # Extract capabilities
        capabilities = []
        for cap in content_div.select('.capabilitiesBoxes'):
            capabilities.append(cap.get_text(strip=True))
        if capabilities:
            result['capabilities'] = capabilities

        return result

    def _extract_statistics_section(self, content_div, description=None, metric=None):
        """Extract statistics from a section with improved parsing"""
        result = {}

        if description:
            result['description'] = description
        if metric:
            result['metric'] = metric

        stats = {}
        
        # First extract main stats containers
        containers = content_div.select("div.specsGenContainers")
        stats.update(self._process_stat_containers(containers))
        
        # Then extract accordion stats if present
        accordion_containers = content_div.select("div.contentSpecs")
        for accordion in accordion_containers:
            accordion_stats = self._process_stat_containers(
                accordion.select("div.specsGenContainers")
            )
            stats.update(accordion_stats)

        if stats:
            result['statistics'] = stats

        return result

    def _process_stat_containers(self, containers):
        """Process individual stat containers and return extracted stats"""
        stats = {}
        
        for container in containers:
            try:
                # Extract title
                title_elem = container.select_one(".textLarge.textYellow.textBold.textShadow")
                if not title_elem:
                    continue
                
                title_text = title_elem.get_text(strip=True).rstrip(':')

                # Extract rank information
                rank_box = container.select_one(".specsRankBox")
                rank_info = None
                if rank_box:
                    white_text = rank_box.select_one(".textWhite")
                    gray_text = rank_box.select_one(".textLtrGray")
                    if white_text and gray_text:
                        rank = white_text.get_text(strip=True)
                        total = gray_text.get_text(strip=True).replace('/', '')
                        rank_info = f"{rank}/{total}"

                # Extract main value
                value_elem = container.select_one(".textLarge .textWhite.textShadow")
                if not value_elem:
                    value_elem = container.select_one(".textWhite.textShadow")  # Fallback
                
                value = None
                if value_elem:
                    value = value_elem.get_text(strip=True)

                # Extract any additional details (like percentages or notes)
                details = []
                detail_elems = container.select(".textLtrGray")
                for elem in detail_elems:
                    detail_text = elem.get_text(strip=True)
                    if detail_text and '/' not in detail_text:  # Avoid duplicate rank info
                        details.append(detail_text)

                # Create stat entry
                stat_entry = {}
                if value:
                    stat_entry['value'] = value
                if rank_info:
                    stat_entry['rank'] = rank_info
                if details:
                    stat_entry['details'] = details

                # Only add non-empty stat entries
                if stat_entry:
                    stats[title_text] = stat_entry

            except Exception as e:
                logger.warning(f"Error processing stat container: {e}")
                continue

        return stats

    def _get_graph_data_from_page(self):
        """Extract graph data from the page"""
        try:
            script = """
            if (typeof powerIndexGraphData !== 'undefined') {
                return powerIndexGraphData;
            }
            return null;
            """
            data = self.driver.execute_script(script)
            if data:
                return data

            # Fallback to parsing HTML
            return self._get_graph_data_from_html(self.driver.page_source)
        except Exception as e:
            logger.error(f"Error extracting graph data: {e}")
            return {}

    def _get_graph_data_from_html(self, html):
        """Fallback method to extract graph data from HTML"""
        soup = BeautifulSoup(html, 'html.parser')
        for script in soup.find_all('script'):
            if script.string and 'data:' in script.string:
                try:
                    data_match = re.search(r"data:\s*(\[.*?\])", script.string)
                    if data_match:
                        values = json.loads(data_match.group(1).replace("'", '"'))
                        labels = ['Manpower', 'Airpower', 'Land Power', 'Naval Power', 'Financials']
                        return dict(zip(labels, values))
                except Exception as e:
                    logger.error(f"Error parsing graph data: {e}")
        return {}

    def save_to_json(self, data, filename="output.json", save_file=True):
        """Save scraped data to a JSON file"""
        if not save_file:
            return data
            
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=4)
            logger.info(f"Data saved to {filename}")
            return True
        except Exception as e:
            logger.error(f"Error saving to JSON: {e}")
            return False

    def get_section_data(self, country_id: str, section_name: str) -> dict:
        """Get data for a specific section of military information"""
        html = self._get_page(country_id)
        if not html:
            return {}
            
        soup = BeautifulSoup(html, 'html.parser')
        
        # Find the section button
        buttons = soup.find_all("button", class_="collapsible")
        for button in buttons:
            header_text = button.get_text(strip=True).rstrip(' [+]').lower()
            normalized_header = header_text.lower().replace("-", " ").strip()
            
            # Check if this is the section we're looking for
            if section_name.lower() in normalized_header:
                content_div = button.find_next_sibling("div", class_="contentSpecs")
                if content_div:
                    # Process the section based on its type
                    if "at a glance" in normalized_header:
                        return self._extract_at_a_glance(content_div)
                    elif normalized_header == "overview":
                        return self._extract_overview(content_div)
                    elif normalized_header == "financials":
                        return self._extract_statistics_section(
                            content_div,
                            description="Financial metrics related to defense spending and economic indicators.",
                            metric="All monetary values presented in United States Dollar (USD$)."
                        )
                    elif normalized_header == "geography":
                        return self._extract_statistics_section(
                            content_div,
                            description="Geographic metrics including land area, coastline, and shared borders.",
                            metric="All distances presented in kilometers (km)."
                        )
                    elif normalized_header == "manpower":
                        return self._extract_statistics_section(
                            content_div,
                            description="The following values detail the maximum, theoretical persons the nation can commit to a war effort."
                        )
                    elif normalized_header == "airpower":
                        return self._extract_statistics_section(
                            content_div,
                            description="Aerial warfare capabilities including aircraft inventory and readiness."
                        )
                    elif normalized_header in ["land forces", "landforces"]:
                        return self._extract_statistics_section(
                            content_div,
                            description="Land warfare capabilities including tanks, armored vehicles, and artillery."
                        )
                    elif normalized_header in ["naval forces", "navalforces"]:
                        naval_stats = self._extract_statistics_section(
                            content_div,
                            description="Naval warfare capabilities including ships and submarines."
                        )
                        # Extract hull class descriptions if available
                        hull_descriptions = {}
                        hull_containers = content_div.select(".hullClassContainers")
                        for hull in hull_containers:
                            hull_type = hull.select_one(".textNormal")
                            hull_desc = hull.select_one(".textSmall1")
                            if hull_type and hull_desc:
                                hull_descriptions[hull_type.get_text(strip=True)] = hull_desc.get_text(strip=True)
                        if hull_descriptions:
                            naval_stats['hull_types'] = hull_descriptions
                        return naval_stats
                    elif "end use" in normalized_header or "end-use" in normalized_header:
                        products = [p.get_text(strip=True) for p in content_div.select('.prodTitleContainer')]
                        return {
                            'description': "Industries that would become stressed in the event of Total War.",
                            'products': products
                        }
                    elif "natural resources" in normalized_header:
                        return self._extract_statistics_section(
                            content_div,
                            description="Natural resources available to the country."
                        )
                    elif normalized_header == "logistics":
                        return self._extract_statistics_section(
                            content_div,
                            description="Logistics capabilities including transportation infrastructure."
                        )
        
        # If section not found, return empty dict
        return {}
