#!/usr/bin/env python3
from __future__ import annotations

import argparse
import importlib
import os
import sys
import time
import threading
from dataclasses import dataclass
from typing import Optional, Tuple, List

from concurrent.futures import ThreadPoolExecutor, as_completed

import requests
import g4f

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

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

DEFAULT_URL = "https://raw.githubusercontent.com/maruf009sultan/g4f-working/refs/heads/main/working/working_results.txt"

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() or "COLAB_GPU" in sys.environ if hasattr(sys, "environ") else True


class C:
    def __init__(self, enable: bool):
        self.enable = enable
        self.R = "\033[0m" if enable else ""
        self.G = "\033[92m" if enable else ""
        self.Y = "\033[93m" if enable else ""
        self.B = "\033[94m" if enable else ""
        self.D = "\033[90m" if enable else ""
        self.RED = "\033[91m" if enable else ""


def trunc(s: str, n: int = 140) -> 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"


# ---------- Core ----------
@dataclass(frozen=True)
class Entry:
    provider_name: str
    model_name: str
    kind: str

    def as_line(self) -> str:
        return f"{self.provider_name}|{self.model_name}|{self.kind}"

    def short(self) -> str:
        # shorter label for progress line
        mn = self.model_name
        if len(mn) > 42:
            mn = "…" + mn[-41:]
        return f"{self.provider_name}|{mn}"
    
    @property
    def db_name(self) -> str:
        """Generate database-friendly name"""
        name = self.model_name.replace('/', '-').replace(' ', '-').lower()
        if self.kind == 'image':
            name += '-img'
        return name


# ---------- 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))


def save_working_model_to_db(entry: Entry, update_existing: bool = False) -> Tuple[bool, str]:
    """Save a working model to the database"""
    db_name = entry.db_name
    
    try:
        existing = get_ai_model_by_name(db_name)
        
        if existing:
            if update_existing or existing.get('status') != 'working':
                success = update_ai_model(existing['id'], status='working')
                return success, 'updated' if success else 'update_failed'
            return True, 'skipped'
        else:
            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',
                capabilities=get_capabilities(entry.model_name, entry.kind),
                max_tokens=estimate_max_tokens(entry.model_name),
                auth_required=False
            )
            return model_id is not None, 'created' if model_id else 'create_failed'
    except Exception as e:
        return False, f'error: {e}'


def fetch_text(url: str, timeout: int = 20) -> str:
    r = requests.get(url, timeout=timeout)
    r.raise_for_status()
    return r.text


def parse_entries(raw: str) -> List[Entry]:
    out: List[Entry] = []
    for line in raw.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        parts = [p.strip() for p in line.split("|")]
        if len(parts) < 3:
            continue
        out.append(Entry(parts[0], parts[1], parts[2]))
    return out


def resolve_provider(provider_name: str):
    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:
    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_good(text: str) -> bool:
    return EXPECTED_TOKEN in (text or "").strip()


def test_one(entry: Entry, timeout: int, proxy: Optional[str]) -> Tuple[bool, str]:
    provider = resolve_provider(entry.provider_name)
    if provider is None:
        return False, f"ProviderNotFound: {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_good(text):
            return True, text
        return False, f"BadResponse: {trunc(text, 200)}"
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"


def main():
    ap = argparse.ArgumentParser(add_help=True)
    ap.add_argument("--url", default=DEFAULT_URL)
    ap.add_argument("--kind", default="text")
    ap.add_argument("--workers", type=int, default=12)
    ap.add_argument("--timeout", type=int, default=20)
    ap.add_argument("--retries", type=int, default=1)
    ap.add_argument("--proxy", default=None)
    ap.add_argument("--fail-log", default="g4f_failures.log")
    ap.add_argument("--no-color", action="store_true", help="Disable ANSI colors")
    ap.add_argument("--summary", action="store_true", help="Print final summary table")
    ap.add_argument("--save-to-db", action="store_true", help="Save working models to database")
    ap.add_argument("--update-existing", action="store_true", help="Update existing models in database")

    # Notebook-safe: ignore unknown args like -f kernel.json
    args, _unknown = ap.parse_known_args()
    
    # Initialize database if saving
    if args.save_to_db:
        print(f"[INFO] Initializing database for saving working models...", file=sys.stderr, flush=True)
        # Note: Database initialization is assumed to be handled elsewhere, e.g., via migrations.
        # if not initialize_ai_models_table():
        #     print(f"[ERROR] Failed to initialize database", file=sys.stderr, flush=True)
        #     return

    color = C(enable=(not args.no_color))

    raw = fetch_text(args.url)
    entries = [e for e in parse_entries(raw) if e.kind.lower() == args.kind.lower()]
    total = len(entries)

    print(f"{color.B}[INFO]{color.R} Loaded {total} '{args.kind}' entries.", file=sys.stderr, flush=True)
    print(f"{color.D}[INFO]{color.R} workers={args.workers} timeout={args.timeout}s retries={args.retries}", file=sys.stderr, flush=True)
    if args.fail_log:
        print(f"{color.D}[INFO]{color.R} failures -> {args.fail_log}", file=sys.stderr, flush=True)

    lock = threading.Lock()
    start_time = time.time()

    done = 0
    ok_count = 0
    failures: List[str] = []
    working: List[str] = []
    
    # Database stats
    db_stats = {
        'created': 0,
        'updated': 0,
        'skipped': 0,
        'failed': 0
    }

    def render_progress(last_label: str, last_status: str):
        # single-line progress update
        elapsed = fmt_elapsed(time.time() - start_time)
        bar_len = 20
        filled = int(bar_len * done / total) if total else bar_len
        bar = "█" * filled + "░" * (bar_len - filled)
        line = (
            f"\r{color.D}[{bar}] {done:>3}/{total:<3}  "
            f"{color.G}OK{color.R}:{ok_count:<3}  "
            f"{color.D}t={elapsed:<8}{color.R}  "
            f"{last_status:<5}  {last_label:<60}"
        )
        # pad to clear leftover chars from previous line
        sys.stderr.write(line + " " * 10)
        sys.stderr.flush()

    def run_with_retries(e: Entry) -> Tuple[Entry, bool, str, int]:
        last_msg = ""
        for attempt in range(1, args.retries + 1):
            ok, msg = test_one(e, timeout=args.timeout, proxy=args.proxy)
            if ok:
                return e, True, msg, attempt
            last_msg = msg
            time.sleep(0.2 * attempt)
        return e, False, last_msg, args.retries

    with ThreadPoolExecutor(max_workers=max(1, args.workers)) as ex:
        futs = [ex.submit(run_with_retries, e) for e in entries]

        for fut in as_completed(futs):
            e, ok, msg, attempts = fut.result()
            with lock:
                done += 1
                last_label = f"{e.short()} (try {attempts})"
                last_status = f"{color.G}OK{color.R}" if ok else f"{color.RED}FAIL{color.R}"
                render_progress(last_label, last_status)

                if ok:
                    line = e.as_line()
                    working.append(line)
                    ok_count += 1
                    
                    # Save to database if enabled
                    db_action = ""
                    if args.save_to_db:
                        db_success, db_action = save_working_model_to_db(e, args.update_existing)
                        if db_action == 'created':
                            db_stats['created'] += 1
                        elif db_action == 'updated':
                            db_stats['updated'] += 1
                        elif db_action == 'skipped':
                            db_stats['skipped'] += 1
                        else:
                            db_stats['failed'] += 1
                    
                    # print a clean success line (new line) without breaking progress
                    sys.stderr.write("\n")
                    db_indicator = f" [{db_action}]" if db_action else ""
                    print(f"{color.G}WORKING{color.R}: {line}{db_indicator}", flush=True)
                else:
                    # keep progress line, then print short error on stderr
                    sys.stderr.write("\n")
                    short_err = trunc(msg, 180)
                    err_line = f"FAILED: {e.as_line()} :: {short_err}"
                    print(f"{color.Y}{err_line}{color.R}", file=sys.stderr, flush=True)
                    failures.append(f"{e.as_line()} :: {msg}")

    # end progress line nicely
    sys.stderr.write("\n")
    sys.stderr.flush()

    # write full failures
    if args.fail_log:
        try:
            with open(args.fail_log, "w", encoding="utf-8") as f:
                for item in failures:
                    f.write(item + "\n")
            print(f"{color.B}[INFO]{color.R} Saved {len(failures)} failures to {args.fail_log}", file=sys.stderr, flush=True)
        except Exception as e:
            print(f"{color.RED}[WARN]{color.R} Could not write fail log: {e}", file=sys.stderr, flush=True)

    # optional summary
    if args.summary:
        print("\n=== SUMMARY ===")
        print(f"Total tested: {total}")
        print(f"Working:      {len(working)}")
        print(f"Failed:       {len(failures)}")
        
        # Database stats
        if args.save_to_db:
            print("\n=== DATABASE OPERATIONS ===")
            print(f"Created:      {db_stats['created']}")
            print(f"Updated:      {db_stats['updated']}")
            print(f"Skipped:      {db_stats['skipped']}")
            print(f"Failed:       {db_stats['failed']}")
        
        if working:
            print("\nWorking list:")
            for w in working:
                print("  " + w)


if __name__ == "__main__":
    main()
