"""
Scrape, Test, and Migrate AI Models from GitHub to Database
Fetches working models from g4f-working repository, tests each one,
and updates database with proper status (working/failed)

Requirements:
    pip install requests g4f

Usage:
    python aimodel_selector/migrate_models_to_db.py [options]

Options:
    --test-only     Only test models, don't update database
    --update-all    Force re-test and update all existing models
    --workers N     Number of parallel workers (default: 8)
    --timeout N     Test timeout in seconds (default: 15)
    --retries N     Number of retries per model (default: 2)
"""

import argparse
import importlib
import requests
import sys
import os
import time
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import Optional, Tuple, List, Dict

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

try:
    import g4f
    G4F_AVAILABLE = True
except ImportError:
    G4F_AVAILABLE = False
    print("⚠️  g4f not installed. Install with: pip install g4f")

from aimodel_selector.db_manager import (
    get_all_ai_models,
    get_ai_model_by_name,
    create_ai_model,
    update_ai_model
)

# GitHub URL with working models
MODELS_URL = "https://raw.githubusercontent.com/maruf009sultan/g4f-working/refs/heads/main/working/working_results.txt"

# Test configuration
EXPECTED_TOKEN = "__G4F_OK_9f3c__"
TEST_PROMPT = f"Reply with exactly this token and nothing else: {EXPECTED_TOKEN}"


# ---------- Pretty printing helpers ----------
def supports_color() -> bool:
    return sys.stderr.isatty()


class Colors:
    def __init__(self, enable: bool = True):
        self.enable = enable
        self.RESET = "\033[0m" if enable else ""
        self.GREEN = "\033[92m" if enable else ""
        self.YELLOW = "\033[93m" if enable else ""
        self.BLUE = "\033[94m" if enable else ""
        self.GRAY = "\033[90m" if enable else ""
        self.RED = "\033[91m" if enable else ""
        self.CYAN = "\033[96m" if enable else ""
        self.BOLD = "\033[1m" if enable else ""


C = Colors(supports_color())


def trunc(s: str, n: int = 100) -> str:
    s = (s or "").replace("\n", " ").replace("\r", " ").strip()
    return s if len(s) <= n else s[: n - 1] + "…"


def fmt_elapsed(sec: float) -> str:
    sec = int(sec)
    m, s = divmod(sec, 60)
    h, m = divmod(m, 60)
    if h:
        return f"{h}h {m:02d}m {s:02d}s"
    if m:
        return f"{m}m {s:02d}s"
    return f"{s}s"


# ---------- Model Data Structures ----------
@dataclass
class ModelEntry:
    provider_name: str
    model_name: str
    capability: str
    
    @property
    def db_name(self) -> str:
        """Generate database-friendly name"""
        name = self.model_name.replace('/', '-').replace(' ', '-').lower()
        if self.capability == 'image':
            name += '-img'
        return name
    
    def short(self) -> str:
        mn = self.model_name
        if len(mn) > 35:
            mn = "…" + mn[-34:]
        return f"{self.provider_name}|{mn}"


# ---------- Testing Functions ----------
def resolve_provider(provider_name: str):
    """Resolve g4f provider by name"""
    if not G4F_AVAILABLE:
        return None
    prov = getattr(g4f.Provider, provider_name, None)
    if prov is not None:
        return prov
    try:
        mod = importlib.import_module("g4f.Provider")
        return getattr(mod, provider_name, None)
    except Exception:
        return None


def normalize_response(resp) -> str:
    """Normalize various response types to string"""
    if resp is None:
        return ""
    if isinstance(resp, str):
        return resp
    if isinstance(resp, dict):
        try:
            return resp["choices"][0]["message"]["content"] or ""
        except Exception:
            return str(resp)
    try:
        chunks = []
        for i, chunk in enumerate(resp):
            if chunk is None:
                continue
            chunks.append(str(chunk))
            if i >= 200:
                break
        return "".join(chunks)
    except Exception:
        return ""


def is_test_passed(text: str) -> bool:
    """Check if response contains expected token"""
    return EXPECTED_TOKEN in (text or "").strip()


def test_model(entry: ModelEntry, timeout: int = 15, proxy: Optional[str] = None) -> Tuple[bool, str]:
    """Test a single model and return (success, message)"""
    if not G4F_AVAILABLE:
        return False, "g4f not available"
    
    provider = resolve_provider(entry.provider_name)
    if provider is None:
        return False, f"Provider not found: {entry.provider_name}"
    
    messages = [{"role": "user", "content": TEST_PROMPT}]
    try:
        g4f.check_version = False
        resp = g4f.ChatCompletion.create(
            model=entry.model_name,
            provider=provider,
            messages=messages,
            timeout=timeout,
            proxy=proxy,
            stream=False,
        )
        text = normalize_response(resp)
        if is_test_passed(text):
            return True, "OK"
        return False, f"Bad response: {trunc(text, 80)}"
    except Exception as e:
        return False, f"{type(e).__name__}: {str(e)[:80]}"


def test_model_with_retries(entry: ModelEntry, timeout: int = 15, retries: int = 2, proxy: Optional[str] = None) -> Tuple[bool, str, int]:
    """Test model with retry logic"""
    last_msg = ""
    for attempt in range(1, retries + 1):
        ok, msg = test_model(entry, timeout=timeout, proxy=proxy)
        if ok:
            return True, msg, attempt
        last_msg = msg
        if attempt < retries:
            time.sleep(0.3 * attempt)
    return False, last_msg, retries


# ---------- Database Helper Functions ----------
def generate_description(provider: str, model_name: str) -> str:
    """Generate a user-friendly description from model name"""
    clean_name = model_name.replace('-', ' ').replace('_', ' ').title()
    
    model_lower = model_name.lower()
    if 'gpt' in model_lower:
        return f"GPT Model via {provider}"
    elif 'claude' in model_lower:
        return f"Claude AI via {provider}"
    elif 'llama' in model_lower:
        return f"Meta Llama via {provider}"
    elif 'mistral' in model_lower:
        return f"Mistral AI via {provider}"
    elif 'qwen' in model_lower:
        return f"Alibaba Qwen via {provider}"
    elif 'deepseek' in model_lower:
        return f"DeepSeek AI via {provider}"
    elif 'gemma' in model_lower:
        return f"Google Gemma via {provider}"
    elif 'flux' in model_lower:
        return f"Flux Image Model via {provider}"
    elif 'command' in model_lower:
        return f"Cohere Command via {provider}"
    elif 'glm' in model_lower:
        return f"GLM Model via {provider}"
    elif 'gemini' in model_lower:
        return f"Google Gemini via {provider}"
    else:
        return f"{clean_name} via {provider}"


def estimate_max_tokens(model_name: str) -> int:
    """Estimate max tokens based on model name"""
    model_lower = model_name.lower()
    
    if '405b' in model_lower or '480b' in model_lower:
        return 128000
    elif '235b' in model_lower or '120b' in model_lower:
        return 64000
    elif '70b' in model_lower or '72b' in model_lower:
        return 32768
    elif '32b' in model_lower or '30b' in model_lower:
        return 16384
    elif '13b' in model_lower or '14b' in model_lower:
        return 8192
    elif '7b' in model_lower or '8b' in model_lower:
        return 4096
    elif '3b' in model_lower or '1b' in model_lower:
        return 2048
    else:
        return 8192


def get_capabilities(model_name: str, base_capability: str) -> List[str]:
    """Determine model capabilities"""
    capabilities = [base_capability]
    model_lower = model_name.lower()
    if 'vision' in model_lower or 'multimodal' in model_lower:
        capabilities.append('vision')
    if 'coder' in model_lower or 'code' in model_lower:
        capabilities.append('code')
    return list(set(capabilities))


# ---------- Core Migration Logic ----------
def fetch_models_from_github() -> str:
    """Fetch models list from GitHub"""
    print(f"{C.BLUE}🌐 Fetching models from GitHub...{C.RESET}")
    response = requests.get(MODELS_URL, timeout=30)
    response.raise_for_status()
    print(f"{C.GREEN}✅ Fetched {len(response.text)} characters{C.RESET}")
    return response.text


def parse_model_entries(raw_text: str) -> List[ModelEntry]:
    """Parse pipe-delimited model data"""
    entries = []
    seen = set()
    
    for line in raw_text.strip().split('\n'):
        line = line.strip()
        if not line or '|' not in line or line.startswith('#'):
            continue
        
        parts = line.split('|')
        if len(parts) != 3:
            continue
        
        provider, model_name, capability = parts
        
        # Create unique identifier
        unique_key = f"{provider}|{model_name}|{capability}"
        if unique_key in seen:
            continue
        seen.add(unique_key)
        
        entries.append(ModelEntry(provider, model_name, capability))
    
    return entries


def migrate_models(
    workers: int = 8,
    timeout: int = 15,
    retries: int = 2,
    test_only: bool = False,
    update_all: bool = False,
    proxy: Optional[str] = None
):
    """Main migration function - test and migrate models"""
    
    print(f"\n{C.BOLD}{C.CYAN}{'='*70}{C.RESET}")
    print(f"{C.BOLD}{C.CYAN}   AI Models Migration & Testing Script{C.RESET}")
    print(f"{C.BOLD}{C.CYAN}{'='*70}{C.RESET}\n")
    
    # Fetch data from GitHub
    try:
        raw_text = fetch_models_from_github()
    except Exception as e:
        print(f"{C.RED}❌ Failed to fetch data: {e}{C.RESET}")
        return False
    
    # Parse entries
    entries = parse_model_entries(raw_text)
    total = len(entries)
    print(f"{C.BLUE}📋 Found {total} unique model entries{C.RESET}")
    
    # Filter only text models for testing (skip image models)
    text_entries = [e for e in entries if e.capability.lower() == 'text']
    image_entries = [e for e in entries if e.capability.lower() == 'image']
    
    print(f"{C.GRAY}   - Text models: {len(text_entries)}{C.RESET}")
    print(f"{C.GRAY}   - Image models: {len(image_entries)} (will be added without testing){C.RESET}")
    
    # Initialize database
    if not test_only:
        print(f"\n{C.BLUE}📋 Initializing database...{C.RESET}")
        # Note: Database initialization is assumed to be handled elsewhere, e.g., via migrations.
        # if not initialize_ai_models_table():
        #     print(f"{C.RED}❌ Failed to initialize database{C.RESET}")
        #     return False
    
    # Stats
    lock = threading.Lock()
    start_time = time.time()
    stats = {
        'tested': 0,
        'working': 0,
        'failed': 0,
        'created': 0,
        'updated': 0,
        'skipped': 0
    }
    results: List[Dict] = []
    
    def process_text_model(entry: ModelEntry) -> Dict:
        """Test and optionally update database for a text model"""
        db_name = entry.db_name
        
        # Check if exists in database
        existing = get_ai_model_by_name(db_name) if not test_only else None
        
        # Test the model
        ok, msg, attempts = test_model_with_retries(entry, timeout=timeout, retries=retries, proxy=proxy)
        status = 'working' if ok else 'inactive'  # Use 'inactive' for failed models (DB ENUM constraint)
        
        result = {
            'entry': entry,
            'db_name': db_name,
            'ok': ok,
            'status': status,
            'message': msg,
            'attempts': attempts,
            'existing': existing is not None,
            'action': 'none'
        }
        
        if test_only:
            return result
        
        # Database operations
        if existing:
            # Update existing model if status changed or update_all flag
            if update_all or existing.get('status') != status:
                success = update_ai_model(existing['id'], status=status)
                result['action'] = 'updated' if success else 'update_failed'
            else:
                result['action'] = 'skipped'
        else:
            # Create new model
            model_id = create_ai_model(
                name=db_name,
                description=generate_description(entry.provider_name, entry.model_name),
                tier='free',
                model=entry.model_name,
                provider=entry.provider_name,
                api_key=None,
                base_url=None,
                status=status,
                capabilities=get_capabilities(entry.model_name, entry.capability),
                max_tokens=estimate_max_tokens(entry.model_name),
                auth_required=False
            )
            result['action'] = 'created' if model_id else 'create_failed'
        
        return result
    
    def render_progress(done: int, total: int, last_label: str, last_ok: bool):
        elapsed = fmt_elapsed(time.time() - start_time)
        bar_len = 25
        filled = int(bar_len * done / total) if total else bar_len
        bar = "█" * filled + "░" * (bar_len - filled)
        status_str = f"{C.GREEN}✓{C.RESET}" if last_ok else f"{C.RED}✗{C.RESET}"
        line = (
            f"\r{C.GRAY}[{bar}]{C.RESET} {done:>3}/{total}  "
            f"{C.GREEN}OK{C.RESET}:{stats['working']:<3} "
            f"{C.RED}OFF{C.RESET}:{stats['failed']:<3}  "
            f"{C.GRAY}t={elapsed:<8}{C.RESET}  "
            f"{status_str} {last_label:<50}"
        )
        sys.stderr.write(line + " " * 5)
        sys.stderr.flush()
    
    # Test text models in parallel
    print(f"\n{C.BLUE}🧪 Testing {len(text_entries)} text models...{C.RESET}")
    print(f"{C.GRAY}   Workers: {workers}, Timeout: {timeout}s, Retries: {retries}{C.RESET}\n")
    
    with ThreadPoolExecutor(max_workers=max(1, workers)) as executor:
        futures = {executor.submit(process_text_model, e): e for e in text_entries}
        
        for future in as_completed(futures):
            result = future.result()
            entry = result['entry']
            
            with lock:
                stats['tested'] += 1
                if result['ok']:
                    stats['working'] += 1
                else:
                    stats['failed'] += 1
                
                if result['action'] == 'created':
                    stats['created'] += 1
                elif result['action'] == 'updated':
                    stats['updated'] += 1
                elif result['action'] == 'skipped':
                    stats['skipped'] += 1
                
                results.append(result)
                render_progress(stats['tested'], len(text_entries), entry.short(), result['ok'])
    
    sys.stderr.write("\n\n")
    
    # Add image models without testing
    if not test_only and image_entries:
        print(f"{C.BLUE}🖼️  Adding {len(image_entries)} image models (without testing)...{C.RESET}")
        for entry in image_entries:
            db_name = entry.db_name
            existing = get_ai_model_by_name(db_name)
            
            if existing:
                stats['skipped'] += 1
                continue
            
            model_id = create_ai_model(
                name=db_name,
                description=generate_description(entry.provider_name, entry.model_name),
                tier='free',
                model=entry.model_name,
                provider=entry.provider_name,
                api_key=None,
                base_url=None,
                status='working',  # Assume working since we can't test image models easily
                capabilities=get_capabilities(entry.model_name, entry.capability),
                max_tokens=estimate_max_tokens(entry.model_name),
                auth_required=False
            )
            if model_id:
                stats['created'] += 1
                print(f"  {C.GREEN}✅{C.RESET} Added: {db_name}")
            else:
                print(f"  {C.RED}❌{C.RESET} Failed: {db_name}")
    
    # Print summary
    elapsed = fmt_elapsed(time.time() - start_time)
    
    print(f"\n{C.BOLD}{'='*70}{C.RESET}")
    print(f"{C.BOLD}   MIGRATION SUMMARY{C.RESET}")
    print(f"{'='*70}")
    print(f"⏱️  Total time: {elapsed}")
    print(f"🧪 Models tested: {stats['tested']}")
    print(f"   {C.GREEN}✓ Working: {stats['working']}{C.RESET}")
    print(f"   {C.RED}✗ Inactive: {stats['failed']}{C.RESET}")
    
    if not test_only:
        print(f"\n📊 Database operations:")
        print(f"   {C.GREEN}➕ Created: {stats['created']}{C.RESET}")
        print(f"   {C.BLUE}🔄 Updated: {stats['updated']}{C.RESET}")
        print(f"   {C.GRAY}⏭️  Skipped: {stats['skipped']}{C.RESET}")
        
        # Verify database
        all_models = get_all_ai_models(active_only=False)
        working_models = [m for m in all_models if m.get('status') == 'working']
        
        print(f"\n📦 Database status:")
        print(f"   Total models: {len(all_models)}")
        print(f"   Working models: {len(working_models)}")
        
        # Provider breakdown
        providers = {}
        for model in working_models:
            provider = model['provider']
            providers[provider] = providers.get(provider, 0) + 1
        
        if providers:
            print(f"\n🏢 Working models by provider:")
            for provider, count in sorted(providers.items(), key=lambda x: x[1], reverse=True)[:10]:
                print(f"   {provider}: {count}")
    
    # Print working models list
    working_results = [r for r in results if r['ok']]
    if working_results:
        print(f"\n{C.GREEN}✅ Working models ({len(working_results)}):{C.RESET}")
        for r in working_results[:20]:  # Show first 20
            entry = r['entry']
            print(f"   {entry.provider_name}|{entry.model_name}")
        if len(working_results) > 20:
            print(f"   ... and {len(working_results) - 20} more")
    
    print(f"\n{'='*70}")
    print(f"{C.GREEN}✅ Migration complete!{C.RESET}")
    print(f"{'='*70}\n")
    
    return True


def main():
    parser = argparse.ArgumentParser(
        description="Scrape, test, and migrate AI models to database"
    )
    parser.add_argument(
        "--test-only", 
        action="store_true",
        help="Only test models, don't update database"
    )
    parser.add_argument(
        "--update-all",
        action="store_true", 
        help="Force re-test and update all existing models"
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=8,
        help="Number of parallel workers (default: 8)"
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=15,
        help="Test timeout in seconds (default: 15)"
    )
    parser.add_argument(
        "--retries",
        type=int,
        default=2,
        help="Number of retries per model (default: 2)"
    )
    parser.add_argument(
        "--proxy",
        default=None,
        help="Proxy URL to use for testing"
    )
    
    args = parser.parse_args()
    
    if not G4F_AVAILABLE:
        print(f"{C.RED}❌ g4f package is required. Install with: pip install g4f{C.RESET}")
        sys.exit(1)
    
    success = migrate_models(
        workers=args.workers,
        timeout=args.timeout,
        retries=args.retries,
        test_only=args.test_only,
        update_all=args.update_all,
        proxy=args.proxy
    )
    
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()
