"""
Zhipu AI model list fetching
"""

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

from .base import FetchBase

# Zhipu model documentation page
ZHIPU_MODEL_DOC_URL = "https://docs.bigmodel.cn/cn/guide/start/model-overview"


class ZhipuModelParser(HTMLParser):
    """HTML Parser for parsing Zhipu model list page"""

    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):
        # Check title to determine current model category
        if tag in ["h2", "h3"]:
            attr_dict = dict(attrs)
            section_id = attr_dict.get("id", "")
            
            # Ignore "Model Overview" and "Recommended Models", focus on specific categories
            if "模型" in section_id:
                if section_id in ["Model Overview", "Recommended Models"]:
                    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
            # We assume model name is in the first column (index 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 in valid category, ignore
                if not self.current_category:
                    return
                    
                name = self.current_model_name.strip()
                # Filter out table headers or invalid content
                if name and name != "模型":
                    if self.current_category not in self.models_by_category:
                        self.models_by_category[self.current_category] = []
                    # Avoid duplicate addition in same 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 FetchZhipu(FetchBase):
    """Zhipu AI model list fetcher"""

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

    def fetch_models(self) -> List[str]:
        """Fetch available model list from Zhipu official documentation page.
        
        Returns:
            List of model ID strings (flattened from all categories)
        """
        try:
            html = self._fetch_html(ZHIPU_MODEL_DOC_URL)
            categorized = self._extract_model_names(html)
            return self._flatten_categories(categorized)
        except urllib.error.URLError as e:
            self._log_error(f"[{self.provider}] HTTP request failed: {e}", exc_info=True)
            return []
        except Exception as e:
            self._log_error(f"[{self.provider}] Unexpected error: {e}", exc_info=True)
            return []

    def _flatten_categories(self, categorized: Dict[str, List[str]]) -> List[str]:
        """Flatten categorized model dict into a single list.
        
        Args:
            categorized: Dict mapping category names to lists of model IDs
            
        Returns:
            Flattened list containing all model IDs
        """
        result = []
        for models in categorized.values():
            result.extend(models)
        return result

    def _fetch_html(self, url: str) -> str:
        """Request specified URL and return HTML content."""
        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:
            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]]:
        """Parse HTML to extract model names, return categorized dictionary"""
        parser = ZhipuModelParser()
        parser.feed(html)
        
        # Return categorized model list
        # Convert model names to lowercase uniformly (optional, if needed)
        result = {}
        for category, models in parser.models_by_category.items():
            result[category] = [m.lower() for m in models]
            
        return result

