#!/usr/bin/env python3
"""测试智谱模型获取器的原始输出"""

from html.parser import HTMLParser
import urllib.request
from typing import List, Set, Dict, Optional, Any
import json

# 直接复制 ZhipuModelFetcher 的代码，避免依赖问题

# 智谱模型文档页面
ZHIPU_MODEL_DOC_URL = "https://docs.bigmodel.cn/cn/guide/start/model-overview"

class ZhipuModelParser(HTMLParser):
    """解析智谱模型列表页面的 HTML Parser"""

    def __init__(self):
        super().__init__()
        self.in_table = False
        self.in_tbody = False
        self.in_tr = False
        self.in_td = False
        self.col_index = -1
        self.current_model_name = ""
        self.capture_text = False
        
        self.current_category: Optional[str] = None
        self.models_by_category: Dict[str, List[str]] = {}

    def handle_starttag(self, tag, attrs):
        # 检查标题，确定当前模型类别
        if tag in ["h2", "h3"]:
            attr_dict = dict(attrs)
            section_id = attr_dict.get("id", "")
            
            # 忽略 "模型一览" 和 "推荐模型"，只关注具体分类
            if "模型" in section_id:
                if section_id in ["模型一览", "推荐模型"]:
                    self.current_category = None
                else:
                    self.current_category = section_id

        if tag == "table":
            self.in_table = True
        elif tag == "tbody":
            self.in_tbody = True
        elif tag == "tr":
            self.in_tr = True
            self.col_index = -1
        elif tag == "td":
            self.in_td = True
            self.col_index += 1
            # 我们假设模型名称在第一列 (索引 0)
            if self.in_table and self.col_index == 0:
                self.capture_text = True
                self.current_model_name = ""

    def handle_endtag(self, tag):
        if tag == "table":
            self.in_table = False
        elif tag == "tbody":
            self.in_tbody = False
        elif tag == "tr":
            self.in_tr = False
        elif tag == "td":
            self.in_td = False
            if self.capture_text:
                self.capture_text = False
                
                # 如果当前不在有效分类中，忽略
                if not self.current_category:
                    return
                    
                name = self.current_model_name.strip()
                # 过滤掉表头或无效内容
                if name and name != "模型":
                    if self.current_category not in self.models_by_category:
                        self.models_by_category[self.current_category] = []
                    # 避免同一分类下重复添加
                    if name not in self.models_by_category[self.current_category]:
                         self.models_by_category[self.current_category].append(name)

    def handle_data(self, data):
        if self.capture_text:
            self.current_model_name += data

class ZhipuModelFetcher:
    """智谱 AI 模型列表获取器"""

    @property
    def provider(self) -> str:
        return "zhipu"

    def fetch_models(self) -> Dict[str, List[str]]:
        """从智谱官方文档页面获取可用模型列表，按类别分组"""
        html = self._fetch_html(ZHIPU_MODEL_DOC_URL)
        return self._extract_model_names(html)

    def _fetch_html(self, url: str) -> str:
        """请求指定 URL 并返回 HTML 内容。"""
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
            "Accept-Encoding": "gzip, deflate, br"
        }
        req = urllib.request.Request(url, headers=headers)
        with urllib.request.urlopen(req, timeout=30) as resp:
            # print(f"DEBUG: Response Headers: {resp.headers}")
            raw = resp.read()
            encoding = resp.headers.get("Content-Encoding")
            if encoding == "gzip":
                import gzip
                raw = gzip.decompress(raw)
            elif encoding == "br":
                import brotli
                raw = brotli.decompress(raw)
        return raw.decode("utf-8", errors="ignore")

    def _extract_model_names(self, html: str) -> Dict[str, List[str]]:
        """解析 HTML 提取模型名称，返回分类字典"""
        parser = ZhipuModelParser()
        parser.feed(html)
        
        # 返回分类后的模型列表
        # 将模型名称统一转小写（可选，如果需要）
        result = {}
        for category, models in parser.models_by_category.items():
            result[category] = [m.lower() for m in models]
            
        return result

def test_zhipu_fetcher():
    fetcher = ZhipuModelFetcher()
    result = fetcher.fetch_models()
    print(result)
    return result

if __name__ == "__main__":
    test_zhipu_fetcher()