from playwright.sync_api import sync_playwright, BrowserContext, Page
import time
import json
import os
from typing import List, Dict, Optional

from company_spider.zhilianzhaopin_company.searcc_kw import generate_url


class CityLoader:
    _instance = None
    
    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super(CityLoader, cls).__new__(cls)
        return cls._instance

    def __init__(self, city_file="city.json"):
        if hasattr(self, 'city_map'):
            return
            
        current_dir = os.path.dirname(os.path.abspath(__file__))
        self.file_path = os.path.join(current_dir, city_file)
        self.city_map = {}
        self._load_cities()

    def _load_cities(self):
        if not os.path.exists(self.file_path):
            print(f"City file not found: {self.file_path}")
            return
        
        try:
            with open(self.file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                self._parse_city_data(data.get("allCity", []))
        except Exception as e:
            print(f"Error loading city file: {e}")

    def _parse_city_data(self, cities):
        for city in cities:
            self.city_map[city['name']] = city['code']
            if 'sublist' in city and city['sublist']:
                self._parse_city_data(city['sublist'])

    def get_code(self, city_name):
        return self.city_map.get(city_name)


def get_companies_from_page(page: Page) -> List[Dict[str, str]]:
    """从搜索结果页面获取公司名称和链接"""
    companies = []
    
    # 尝试多种选择器来定位公司名称
    company_selectors = [
        'a[class*="company"]',
        '.company-name a',
        'a.company-name',
        '[class*="CompanyName"] a',
        'a[href*="/company/"]'
    ]
    
    company_elements = []
    for selector in company_selectors:
        try:
            elements = page.query_selector_all(selector)
            if elements:
                company_elements = elements
                print(f"使用选择器找到 {len(elements)} 个元素: {selector}")
                break
        except:
            continue
    
    # 如果没找到，尝试更通用的方法
    if not company_elements:
        all_links = page.query_selector_all('a[href*="company"]')
        company_elements = all_links
        print(f"通过通用方法找到 {len(all_links)} 个公司链接")
    
    # 提取公司信息
    company_info_set = set()
    
    for element in company_elements:
        try:
            company_name = element.inner_text().strip()
            company_url = element.get_attribute('href')
            
            if company_name and company_url:
                # 处理相对路径
                if company_url.startswith('/'):
                    company_url = f"https://www.zhaopin.com{company_url}"
                elif not company_url.startswith('http'):
                    company_url = f"https://www.zhaopin.com/{company_url}"
                
                # 去重
                if company_name not in company_info_set:
                    company_info_set.add(company_name)
                    companies.append({
                        'name': company_name,
                        'url': company_url
                    })
        except Exception as e:
            continue
    
    return companies


def get_company_intro(context: BrowserContext, company_url: str) -> str:
    """获取公司详情简介"""
    try:
        company_page = context.new_page()
        company_page.goto(company_url, wait_until="networkidle", timeout=30000)
        time.sleep(2)
        
        # 尝试多种选择器获取公司简介
        intro_selectors = [
            '.company-intro',
            '.company-description',
            '[class*="intro"]',
            '[class*="description"]',
            '.company-info',
            '[class*="CompanyIntro"]'
        ]
        
        company_intro = ""
        for selector in intro_selectors:
            try:
                intro_element = company_page.query_selector(selector)
                if intro_element:
                    company_intro = intro_element.inner_text().strip()
                    if company_intro:
                        break
            except:
                continue
        
        # 如果还是没找到，尝试获取页面主要内容
        if not company_intro:
            try:
                body = company_page.query_selector('body')
                if body:
                    all_text = body.inner_text()
                    company_intro = all_text[:500]
            except:
                pass
        
        company_page.close()
        return company_intro if company_intro else "未找到公司简介"
        
    except Exception as e:
        return f"获取失败: {str(e)}"


def crawl_companies(params: Dict, max_companies: int = 10, headless: bool = False, proxy: Optional[str] = None) -> List[Dict]:
    """
    爬取智联招聘公司信息
    
    Args:
        params: 搜索参数，如 {'jl': 530, 'kw': 'app推广经理'} 或 {'city': '北京', 'kw': '...'}
        max_companies: 最多爬取的公司数量，默认10
        headless: 是否无头模式，默认False
        proxy: 代理地址，例如 "http://user:pass@host:port"
    
    Returns:
        公司信息列表，每个元素包含 name, url, intro
        如果找到完全匹配的公司名称，只返回该公司的信息（列表长度为1）
    """
    # 处理城市名称转代码
    if 'city' in params and 'jl' not in params:
        city_loader = CityLoader()
        code = city_loader.get_code(params['city'])
        if code:
            print(f"城市 '{params['city']}' 映射代码为: {code}")
            params['jl'] = code
        else:
            print(f"未找到城市 '{params['city']}' 的代码")

    target_company = params.get('kw', '').strip()
    
    with sync_playwright() as p:
        launch_args = ["--disable-blink-features=AutomationControlled"]
        browser_kwargs = {
            "headless": headless,
            "args": launch_args
        }
        
        # 尝试使用本地Chrome，如果不存在则使用默认浏览器
        chrome_path = r"C:\Program Files\Google\Chrome\Application\chrome.exe"
        if os.path.exists(chrome_path):
            browser_kwargs["executable_path"] = chrome_path
        
        if proxy:
            browser_kwargs["proxy"] = {"server": proxy}
            print(f"使用代理: {proxy}")

        browser = p.chromium.launch(**browser_kwargs)
        
        context = browser.new_context(
            viewport={"width": 1920, "height": 1080},
            user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
        )
        
        page = context.new_page()
        
        # 生成URL并访问搜索页面
        url = f"https://www.zhaopin.com/sou/{generate_url(params)}"
        print(f"访问URL: {url}")
        page.goto(url, wait_until="networkidle", timeout=30000)
        time.sleep(3)
        
        # 获取公司列表
        companies = get_companies_from_page(page)
        print(f"\n找到 {len(companies)} 家公司")
        
        # 如果有关键词，尝试精确匹配公司名称
        if target_company:
            print(f"搜索目标公司: {target_company}")
            for company in companies:
                company_name = company['name'].strip()
                # 优先精确匹配，如果精确匹配失败则尝试包含匹配
                if company_name == target_company:
                    print(f"找到完全匹配的公司: {company_name}")
                    print(f"正在获取公司简介...")
                    company_intro = get_company_intro(context, company['url'])
                    
                    context.close()
                    browser.close()
                    
                    return [{
                        'name': company_name,
                        'url': company['url'],
                        'intro': company_intro
                    }]
            
            # 如果精确匹配失败，尝试包含匹配
            for company in companies:
                company_name = company['name'].strip()
                if target_company in company_name or company_name in target_company:
                    print(f"找到部分匹配的公司: {company_name}")
                    print(f"正在获取公司简介...")
                    company_intro = get_company_intro(context, company['url'])
                    
                    context.close()
                    browser.close()
                    
                    return [{
                        'name': company_name,
                        'url': company['url'],
                        'intro': company_intro
                    }]
        
        # 如果没有找到匹配的公司，按原逻辑获取多家公司
        print(f"未找到完全匹配的公司，获取前 {max_companies} 家公司信息")
        results = []
        for i, company in enumerate(companies[:max_companies], 1):
            print(f"\n[{i}/{min(max_companies, len(companies))}] 正在获取: {company['name']}")
            company_intro = get_company_intro(context, company['url'])
            
            results.append({
                'name': company['name'],
                'url': company['url'],
                'intro': company_intro
            })
            
            time.sleep(1)
        
        context.close()
        browser.close()
        
        return results


if __name__ == '__main__':
    # 测试代码
    params = {'city': '北京', 'kw': 'app推广经理'}
    results = crawl_companies(params, max_companies=10)
    
    # 输出结果
    print("\n" + "="*80)
    print("爬取结果:")
    print("="*80)
    for result in results:
        print(f"\n公司名称: {result['name']}")
        print(f"公司链接: {result['url']}")
        print(f"公司简介: {result['intro'][:200]}..." if len(result['intro']) > 200 else f"公司简介: {result['intro']}")
        print("-"*80)
