#!/usr/bin/env python3
"""DrissionPage批量抓取成都区县组织部/党校2026年预算"""
import json, time, os, re, traceback
from DrissionPage import ChromiumPage, ChromiumOptions

OUT_DIR = "/data/www/files/budget_pdfs/chengdu_districts"
RESULT_FILE = "/data/www/files/budget_data/cd_districts_results.json"
os.makedirs(OUT_DIR, exist_ok=True)
os.makedirs(os.path.dirname(RESULT_FILE), exist_ok=True)

# 全部20个区县配置
DISTRICTS = {
    '金牛区': {'domain': 'www.jinniu.gov.cn', 'code': 'jinniu', 'type': '区'},
    '锦江区': {'domain': 'www.jinjiang.gov.cn', 'code': 'jinjiang', 'type': '区'},
    '青羊区': {'domain': 'www.qingyang.gov.cn', 'code': 'qingyang', 'type': '区'},
    '武侯区': {'domain': 'www.wuhou.gov.cn', 'code': 'wuhou', 'type': '区'},
    '成华区': {'domain': 'www.chenghua.gov.cn', 'code': 'chenghua', 'type': '区'},
    '新都区': {'domain': 'www.xindu.gov.cn', 'code': 'xindu', 'type': '区'},
    '温江区': {'domain': 'www.wenjiang.gov.cn', 'code': 'wenjiang', 'type': '区'},
    '双流区': {'domain': 'www.shuangliu.gov.cn', 'code': 'shuangliu', 'type': '区'},
    '郫都区': {'domain': 'www.pidu.gov.cn', 'code': 'pidu', 'type': '区'},
    '新津区': {'domain': 'www.xinjin.gov.cn', 'code': 'xinjin', 'type': '区'},
    '青白江区': {'domain': 'www.qbj.gov.cn', 'code': 'qingbaijiang', 'type': '区'},
    '龙泉驿区': {'domain': 'www.longquanyi.gov.cn', 'code': 'longquanyi', 'type': '区'},
    '都江堰市': {'domain': 'www.djy.gov.cn', 'code': 'dujiangyan', 'type': '市'},
    '彭州市': {'domain': 'www.pengzhou.gov.cn', 'code': 'pengzhou', 'type': '市'},
    '邛崃市': {'domain': 'www.qionglai.gov.cn', 'code': 'qionglai', 'type': '市'},
    '崇州市': {'domain': 'www.chongzhou.gov.cn', 'code': 'chongzhou', 'type': '市'},
    '简阳市': {'domain': 'www.jianyang.gov.cn', 'code': 'jianyang', 'type': '市'},
    '金堂县': {'domain': 'www.jintang.gov.cn', 'code': 'jintang', 'type': '县'},
    '大邑县': {'domain': 'www.dayi.gov.cn', 'code': 'dayi', 'type': '县'},
    '蒲江县': {'domain': 'www.pujiang.gov.cn', 'code': 'pujiang', 'type': '县'},
}

# 已覆盖的区县跳过
COVERED = ['金牛区', '锦江区', '双流区', '龙泉驿区']

# 组织部/党校关键字
TARGET_KEYWORDS = ['组织部', '党校']
YEAR = '2026'

results = {}

def init_browser():
    co = ChromiumOptions()
    co.set_argument('--no-sandbox')
    co.set_argument('--disable-dev-shm-usage')
    co.set_argument('--disable-gpu')
    co.set_argument('--window-size=1920,1080')
    co.headless(True)
    page = ChromiumPage(co)
    return page

def find_budget_links(page, district_name, cfg):
    """在区县财政局'区级部门财政预决算'栏目中找组织部和党校的预算链接"""
    domain = cfg['domain']
    code = cfg['code']
    district_type = cfg['type']
    
    found = {'组织部': [], '党校': []}
    
    # 方法1：直接访问财政信息栏目，搜索预决算公开
    urls_to_try = [
        f"https://{domain}/{code}/c110728/list.shtml",  # 财政信息
        f"https://{domain}/{code}/c110728/2026_cztz.shtml",  # 财政通知
        f"https://{domain}/{code}/c110729/list.shtml",  # 预决算
    ]
    
    for base_url in urls_to_try:
        try:
            page.get(base_url)
            time.sleep(2)
            
            # 搜索页面中的组织部/党校链接
            links = page.eles('tag:a')
            for link in links:
                try:
                    href = link.attr('href') or ''
                    text = link.text or ''
                    for kw in TARGET_KEYWORDS:
                        if kw in text and YEAR in text:
                            full_href = href if href.startswith('http') else f"https://{domain}{href}" if href.startswith('/') else f"https://{domain}/{code}/{href}"
                            title = text.strip()
                            if title not in [f['title'] for f in found[kw]]:
                                found[kw].append({'title': title, 'url': full_href})
                except:
                    continue
        except Exception as e:
            continue
    
    # 方法2：用站内搜索功能
    search_url = f"https://{domain}/search.html?searchWord={district_type}委组织部+{YEAR}+预算"
    try:
        page.get(search_url)
        time.sleep(2)
        links = page.eles('tag:a')
        for link in links:
            try:
                href = link.attr('href') or ''
                text = link.text or ''
                if '组织部' in text and '预算' in text:
                    full_href = href if href.startswith('http') else f"https://{domain}{href}"
                    if text.strip() not in [f['title'] for f in found['组织部']]:
                        found['组织部'].append({'title': text.strip(), 'url': full_href})
            except:
                continue
    except:
        pass
    
    # 方法3：直接构造常见URL模式
    patterns = [
        # 组织部
        f"https://{domain}/{code}/c110728/2026/list.shtml",
        # 财政预决算公开专栏
        f"https://{domain}/{code}/c112103/list.shtml",
        f"https://{domain}/{code}/c115761/list.shtml",
    ]
    for p_url in patterns:
        try:
            page.get(p_url)
            time.sleep(1.5)
            html = page.html or ''
            for kw in TARGET_KEYWORDS:
                # 从HTML中提取含关键字的链接
                a_pattern = rf'<a[^>]*href=["\']([^"\']*)["\'][^>]*>[^<]*{kw}[^<]*{YEAR}[^<]*</a>'
                for m in re.finditer(a_pattern, html, re.IGNORECASE):
                    href = m.group(1)
                    full_href = href if href.startswith('http') else f"https://{domain}{href}" if href.startswith('/') else f"https://{domain}/{code}/{href}"
                    text = re.sub(r'<[^>]+>', '', m.group(0)).strip()
                    if text not in [f['title'] for f in found[kw]]:
                        found[kw].append({'title': text, 'url': full_href})
        except:
            continue
    
    # 方法4：遍历财政信息栏目分页列表
    for page_num in range(1, 6):
        list_url = f"https://{domain}/{code}/c110728/list_{page_num}.shtml"
        try:
            page.get(list_url)
            time.sleep(1.5)
            html = page.html or ''
            for kw in TARGET_KEYWORDS:
                a_pattern = rf'<a[^>]*href=["\']([^"\']*)["\'][^>]*>[^<]*{kw}[^<]*</a>'
                for m in re.finditer(a_pattern, html, re.IGNORECASE):
                    href = m.group(1)
                    full_href = href if href.startswith('http') else f"https://{domain}{href}" if href.startswith('/') else f"https://{domain}/{code}/{href}"
                    text = re.sub(r'<[^>]+>', '', m.group(0)).strip()
                    if YEAR in text or '2026' in href:
                        if text not in [f['title'] for f in found[kw]]:
                            found[kw].append({'title': text, 'url': full_href})
        except:
            continue
    
    return found

def download_pdf(page, url, save_path):
    """用浏览器Fetch API下载PDF（绕过WAF）"""
    try:
        # 先访问URL建立WAF Cookie
        page.get(url)
        time.sleep(2)
        
        # 获取当前页面域名
        from urllib.parse import urlparse
        parsed = urlparse(url)
        origin = f"{parsed.scheme}://{parsed.netloc}"
        
        # 用Fetch API下载
        js_code = f"""
        return await fetch('{url}')
            .then(r => r.arrayBuffer())
            .then(buf => Array.from(new Uint8Array(buf)))
        """
        byte_list = page.run_js(js_code)
        if byte_list and len(byte_list) > 100:
            with open(save_path, 'wb') as f:
                f.write(bytes(byte_list))
            return True
    except:
        pass
    
    # 备用方案：直接获取响应内容
    try:
        page.get(url)
        time.sleep(3)
        # 检查是否是PDF
        content = page.run_js("return document.contentType")
        if content and 'pdf' in content.lower():
            js_code = """
            return await fetch(window.location.href)
                .then(r => r.arrayBuffer())
                .then(buf => Array.from(new Uint8Array(buf)))
            """
            byte_list = page.run_js(js_code)
            if byte_list and len(byte_list) > 100:
                with open(save_path, 'wb') as f:
                    f.write(bytes(byte_list))
                return True
    except:
        pass
    
    return False

def extract_attachments(page, content_url, district_name, kw):
    """从预算公开详情页提取附件PDF链接"""
    attachments = []
    try:
        page.get(content_url)
        time.sleep(2)
        
        # 查找附件下载链接
        # 常见模式：.xwb-attachment, .file-list, a[href*='.pdf'], a[href*='.xlsx']
        file_links = page.eles('css:a[href*=".pdf"], css:a[href*=".xlsx"], css:a[href*=".xls"], css:a[href*=".doc"]')
        for fl in file_links:
            try:
                href = fl.attr('href') or ''
                text = fl.text or ''
                if not href:
                    continue
                # 转绝对URL
                if href.startswith('/'):
                    from urllib.parse import urlparse
                    parsed = urlparse(content_url)
                    href = f"{parsed.scheme}://{parsed.netloc}{href}"
                elif not href.startswith('http'):
                    parsed = urlparse(content_url)
                    path = parsed.path.rsplit('/', 1)[0]
                    href = f"{parsed.scheme}://{parsed.netloc}{path}/{href}"
                
                attachments.append({'title': text.strip() or os.path.basename(href), 'url': href})
            except:
                continue
    except:
        pass
    
    return attachments

def take_screenshot(page, url, save_path):
    """截取预算页面截图"""
    try:
        page.get(url)
        time.sleep(2)
        page.get_screenshot(path=save_path, full_page=True)
        return True
    except:
        return False


# ========== 主流程 ==========
print("=" * 60)
print("DrissionPage 批量抓取成都区县组织部/党校2026年预算")
print("=" * 60)

page = init_browser()
print("浏览器启动成功")

for district_name, cfg in DISTRICTS.items():
    if district_name in COVERED:
        print(f"\n⏩ {district_name} 已覆盖，跳过")
        continue
    
    print(f"\n{'='*40}")
    print(f"🔍 正在扫描: {district_name} ({cfg['domain']})")
    print(f"{'='*40}")
    
    district_result = {'组织部门': [], '党校': [], 'screenshots': [], 'status': 'scanning'}
    
    # 访问首页建立WAF Cookie
    try:
        page.get(f"https://{cfg['domain']}/")
        time.sleep(3)
        title = page.title or ''
        print(f"  首页标题: {title[:50]}")
        
        if '禁止' in title or '403' in title or 'error' in title.lower():
            print(f"  ❌ WAF拦截，跳过")
            district_result['status'] = 'waf_blocked'
            results[district_name] = district_result
            continue
    except Exception as e:
        print(f"  ❌ 无法访问: {e}")
        district_result['status'] = 'unreachable'
        results[district_name] = district_result
        continue
    
    # 查找预算链接
    found = find_budget_links(page, district_name, cfg)
    
    for kw in ['组织部', '党校']:
        kw_key = '组织部门' if kw == '组织部' else '党校'
        print(f"\n  📂 {kw}: 找到 {len(found[kw])} 条链接")
        
        for item in found[kw][:5]:  # 每类限制5条
            print(f"    - {item['title'][:60]}")
            print(f"      {item['url'][:80]}")
            
            # 截图
            ss_name = f"{cfg['code']}_{kw}_{YEAR}.png".replace(' ', '')
            ss_path = f"{OUT_DIR}/../screenshots/{ss_name}"
            os.makedirs(os.path.dirname(ss_path), exist_ok=True)
            if take_screenshot(page, item['url'], ss_path):
                district_result['screenshots'].append({
                    'file': f"budget_pdfs/screenshots/{ss_name}",
                    'url': item['url'],
                    'title': item['title']
                })
            
            # 提取附件
            attachments = extract_attachments(page, item['url'], district_name, kw)
            print(f"      附件: {len(attachments)}个")
            
            entry = {
                'title': item['title'],
                'source_url': item['url'],
                'attachments': []
            }
            
            for att in attachments[:3]:  # 每条限制3个附件
                att_name = f"{cfg['code']}_{kw}_{att['title'][:20]}"
                att_name = re.sub(r'[^\w\u4e00-\u9fff.]', '_', att_name)
                att_path = f"{OUT_DIR}/{att_name}"
                
                if download_pdf(page, att['url'], att_path):
                    size = os.path.getsize(att_path)
                    print(f"      ✅ 下载: {att_name} ({size/1024:.0f}KB)")
                    entry['attachments'].append({
                        'title': att['title'],
                        'url': att['url'],
                        'local_file': f"budget_pdfs/chengdu_districts/{att_name}",
                        'size': size
                    })
                else:
                    print(f"      ❌ 下载失败: {att_name}")
            
            district_result[kw_key].append(entry)
    
    # 如果方法1-4都没找到，尝试直接搜索财政预决算公开专栏
    if not district_result['组织部门'] and not district_result['党校']:
        print(f"  🔄 尝试直接搜索财政预决算专栏...")
        # 尝试更多URL模式
        extra_urls = [
            f"https://{cfg['domain']}/{cfg['code']}/c110728/{YEAR}/list.shtml",
            f"https://{cfg['domain']}/{cfg['code']}/c100043/{YEAR}/list.shtml",
            f"https://{cfg['domain']}/{cfg['code']}/col/col110728/index.html",
        ]
        for eu in extra_urls:
            try:
                page.get(eu)
                time.sleep(1.5)
                html = page.html or ''
                for kw in TARGET_KEYWORDS:
                    if kw in html and YEAR in html:
                        kw_key = '组织部门' if kw == '组织部' else '党校'
                        # 提取链接
                        links = page.eles('tag:a')
                        for link in links:
                            text = link.text or ''
                            if kw in text and (YEAR in text or '预算' in text):
                                href = link.attr('href') or ''
                                if href.startswith('/'):
                                    href = f"https://{cfg['domain']}{href}"
                                elif not href.startswith('http'):
                                    href = f"https://{cfg['domain']}/{cfg['code']}/{href}"
                                district_result[kw_key].append({
                                    'title': text.strip(),
                                    'source_url': href,
                                    'attachments': []
                                })
                                print(f"    ✅ 找到: {text.strip()[:50]}")
            except:
                continue
    
    district_result['status'] = 'completed'
    results[district_name] = district_result
    
    # 每完成一个区县就保存
    with open(RESULT_FILE, 'w') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    print(f"  💾 已保存进度")

# 关闭浏览器
try:
    page.quit()
except:
    pass

# 统计
total_org = sum(len(r.get('组织部门', [])) for r in results.values())
total_dx = sum(len(r.get('党校', [])) for r in results.values())
total_files = sum(
    len(a.get('attachments', [])) 
    for r in results.values() 
    for dept in ['组织部门', '党校'] 
    for a in r.get(dept, [])
)
completed = sum(1 for r in results.values() if r.get('status') == 'completed')

print(f"\n{'='*60}")
print(f"抓取完成统计:")
print(f"  扫描区县: {len(results)}")
print(f"  成功完成: {completed}")
print(f"  组织部条目: {total_org}")
print(f"  党校条目: {total_dx}")
print(f"  下载文件: {total_files}")
print(f"结果文件: {RESULT_FILE}")
print(f"{'='*60}")