import logging
import re
from typing import Dict, List, Optional

import requests
from bs4 import BeautifulSoup

logger = logging.getLogger(__name__)

class OllamaModelValidator:
    """Ollama model validation service"""
    
    def __init__(self):
        self.base_url = "https://ollama.com"
        self.session = requests.Session()
        self.session.headers.update({
            "User-Agent": "Turkish-MMLU-Benchmark/1.0",
            "Accept": "text/html",
        })
        self.session.timeout = 30  # Increased timeout for HF Spaces
    
    def fetch_model_html(self, model_name: str, version: str = None) -> str:
        """Fetch model HTML from Ollama library"""
        # some models https://ollama.com/alibayram/gemma3n
        # some models https://ollama.com/library/gemma3n
        max_retries = 3
        for attempt in range(max_retries):
            try:
                if "/" in model_name:
                    # remove the first /
                    path = f"/{model_name.lstrip('/')}"
                else:
                    path = f"/library/{model_name}"
                if version:
                    path += f":{version}"
                
                url = f"{self.base_url}{path}"
                logger.info(f"Fetching model HTML from: {url} (attempt {attempt + 1}/{max_retries})")
                
                response = self.session.get(url)
                
                if response.status_code == 200:
                    logger.info(f"Successfully fetched HTML from {url}")
                    return response.text
                elif response.status_code == 404:
                    logger.warning(f"Model not found: {url}")
                    raise Exception(f"Model not found: {model_name}:{version}")
                else:
                    logger.warning(f"HTTP {response.status_code} for {url}")
                    raise Exception(f"Failed to fetch model page: {response.status_code}")
                    
            except requests.exceptions.Timeout:
                logger.warning(f"Timeout on attempt {attempt + 1} for {url}")
                if attempt == max_retries - 1:
                    raise Exception(f"Request timeout after {max_retries} attempts")
                continue
            except requests.exceptions.ConnectionError as e:
                logger.warning(f"Connection error on attempt {attempt + 1}: {e}")
                if attempt == max_retries - 1:
                    raise Exception(f"Connection failed after {max_retries} attempts: {e}")
                continue
            except Exception as e:
                logger.error(f"Error fetching model HTML: {e}")
                raise
    
    def parse_model_html(self, html: str) -> Dict:
        """Parse model HTML to extract parameter size and available tags"""
        try:
            soup = BeautifulSoup(html, 'html.parser')
            
            parameter_size = None
            tags = []
            
            # Extract parameter size from body text
            """ <div class="flex sm:space-x-2 items-center"><span class="hidden sm:block">parameters</span><span class="text-neutral-400 sm:font-semibold sm:text-neutral-800 sm:text-xs">6.87B</span></div> """
            body_text = soup.get_text()
            parameter_size_match = re.search(r'parameters\s*(\d+\.\d+B)', body_text)
            if parameter_size_match:
                parameter_size = float(parameter_size_match.group(1).replace('B', ''))
            else:
                logger.warning(f"Could not find parameter size in HTML")
            
            # Extract tags from links
            for link in soup.find_all('a', href=True):
                href = link['href']
                # check if href contains a slash
                if "/" in href:
                    tag_match = re.search(r':([^/]+)', href)
                    if tag_match:
                        tag = tag_match.group(1)
                        if tag not in tags:
                            tags.append(tag)
                if href.startswith('/library/'):
                    tag_match = re.search(r':([^/]+)', href)
                    if tag_match:
                        tag = tag_match.group(1)
                        if tag not in tags:
                            tags.append(tag)
            
            # Extract tags from body text
            tag_matches = re.findall(r'[a-zA-Z0-9_-]+:[a-zA-Z0-9._-]+', body_text)
            for match in tag_matches:
                parts = match.split(':')
                if len(parts) == 2:
                    tag = parts[1]
                    if tag not in tags:
                        tags.append(tag)
            
            # Determine latest tag
            latest_tag = None
            if 'latest' in tags:
                latest_tag = 'latest'
            elif tags:
                latest_tag = tags[0]
            
            return {
                'parameter_size': parameter_size,
                'tags': tags,
                'latest_tag': latest_tag,
            }
            
        except Exception as e:
            logger.error(f"Error parsing model HTML: {e}")
            raise
    
    def calculate_ram_requirement(self, parameter_size: int, quantization_level: int = 8) -> float:
        """Calculate RAM requirements for a model"""
        if quantization_level == 8:
            ram_gb = parameter_size * 1
        elif quantization_level == 16:
            ram_gb = parameter_size * 2
        elif quantization_level == 4:
            ram_gb = parameter_size * 0.5
        else:
            ram_gb = parameter_size * (quantization_level / 8)
        
        return round(ram_gb * 100) / 100
    
    def validate_ollama_model(self, model_name: str, version: str) -> Dict:
        """Validate Ollama model by HTML scraping"""
        try:
            logger.info(f"Validating model: {model_name}:{version or '(latest)'}")
            
            # First, try to fetch the base model to get available tags
            try:
                base_html = self.fetch_model_html(model_name)
                base_parsed = self.parse_model_html(base_html)
                available_tags = base_parsed['tags']
            except Exception as e:
                logger.error(f"Failed to fetch base model {model_name}: {e}")
                return {
                    'valid': False,
                    'error': f"Model {model_name} not found in Ollama library",
                    'available_tags': [],
                }
            
            # If a specific version was requested, check if it exists
            if version:
                if version not in available_tags:
                    return {
                        'valid': False,
                        'error': f"Version '{version}' not found for model {model_name}",
                        'available_tags': available_tags,
                        'suggested_version': base_parsed.get('latest_tag', available_tags[0] if available_tags else None)
                    }
                
                # Version exists, fetch the specific version
                try:
                    html = self.fetch_model_html(model_name, version)
                    parsed_data = self.parse_model_html(html)
                    used_version = version
                except Exception as e:
                    logger.error(f"Failed to fetch specific version {model_name}:{version}: {e}")
                    return {
                        'valid': False,
                        'error': f"Failed to fetch version '{version}' for model {model_name}",
                        'available_tags': available_tags,
                    }
            else:
                # No version specified, use latest
                html = base_html
                parsed_data = base_parsed
                used_version = base_parsed.get('latest_tag', 'latest')
            
            parameter_size = parsed_data['parameter_size']
            
            if not parameter_size:
                return {
                    'valid': False,
                    'error': f"Could not determine parameter size for model {model_name}:{used_version}",
                    'available_tags': available_tags,
                }
            
            # Calculate RAM requirements
            ram_requirements = {
                "16-bit": self.calculate_ram_requirement(parameter_size, 16),
                "8-bit": self.calculate_ram_requirement(parameter_size, 8),
                "4-bit": self.calculate_ram_requirement(parameter_size, 4),
            }
            
            max_ram = 20  # Reverted back to 20GB limit
            fits_in_ram = {
                "16-bit": ram_requirements["16-bit"] <= max_ram,
                "8-bit": ram_requirements["8-bit"] <= max_ram,
                "4-bit": ram_requirements["4-bit"] <= max_ram,
            }
            
            can_load = any(fits_in_ram.values())
            recommended_quantization = None
            
            if can_load:
                for quantization, fits in fits_in_ram.items():
                    if fits:
                        recommended_quantization = quantization
                        break
            
            return {
                'valid': True,
                'model': model_name,
                'version': used_version,
                'parameter_size': parameter_size,
                'ram_requirements': ram_requirements,
                'fits_in_ram': fits_in_ram,
                'can_load': can_load,
                'recommended_quantization': recommended_quantization,
                'available_tags': available_tags,
            }
            
        except Exception as error:
            logger.error(f"Error validating model {model_name}:{version}: {error}")
            return {
                'valid': False,
                'error': f"Validation failed: {str(error)}",
            }

# Global validator instance
model_validator = OllamaModelValidator()