import os, time

def query_model(prompt, max_tokens=1500):
    """Single Groq API call."""
    GROQ_API_KEY = os.getenv("GROQ_API_KEY", "").strip()

    if not GROQ_API_KEY:
        raise ValueError(
            "GROQ_API_KEY is not set.\n\n"
            "Steps to fix:\n"
            "1. Go to https://console.groq.com → sign up free\n"
            "2. Click 'API Keys' → 'Create API Key'\n"
            "3. Copy the key (starts with gsk_...)\n"
            "4. In your HuggingFace Space → Settings → Variables and Secrets\n"
            "   → Add secret: Name = GROQ_API_KEY, Value = gsk_xxxxxxxxxxxx\n"
            "5. Restart the Space"
        )

    try:
        from groq import Groq
        client = Groq(api_key=GROQ_API_KEY)

        response = client.chat.completions.create(
            model="llama-3.3-70b-versatile",   # fastest + best quality on Groq
            messages=[
                {
                    "role": "system",
                    "content": (
                        "You are a certified professional fitness trainer. "
                        "Provide detailed, structured, day-by-day workout plans exactly as instructed."
                    )
                },
                {"role": "user", "content": prompt}
            ],
            max_tokens=max_tokens,
            temperature=0.7,
        )
        return response.choices[0].message.content

    except Exception as e:
        err = str(e)

        if "401" in err or "invalid_api_key" in err.lower() or "authentication" in err.lower():
            raise ValueError(
                "❌ Groq API key is invalid or expired.\n\n"
                "Fix: Go to https://console.groq.com/keys → create a new key → "
                "update GROQ_API_KEY in your HuggingFace Space Secrets."
            ) from None

        if "429" in err or "rate_limit" in err.lower() or "rate limit" in err.lower():
            raise ValueError(
                "⏳ Groq rate limit hit. Wait 60 seconds and try again.\n\n"
                "Free tier: 6,000 tokens/min. Upgrade at console.groq.com for higher limits."
            ) from None

        if "503" in err or "504" in err or "timeout" in err.lower():
            raise ValueError(
                "⚡ Groq server busy. Please try again in a few seconds."
            ) from None

        raise


def query_model_chunked(name, gender, height, weight, goal, fitness_level,
                         equipment, days_per_week=5, months=1,
                         progress_callback=None):
    """
    Generate a long plan by splitting into chunks of 3 days each.
    Groq is fast enough (~2-3s per chunk) so even 30-day plans finish in ~30s.
    progress_callback(chunk_num, total_chunks, days_done, total_days) — optional.
    Returns: (full_plan_text, bmi, bmi_cat)
    """
    from prompt_builder import calculate_bmi, bmi_category, bmi_advice

    bmi      = calculate_bmi(weight, height)
    bmi_cat  = bmi_category(bmi)
    advice   = bmi_advice(bmi_cat)

    total_days = min(days_per_week * 4 * months, 30)
    eq_list    = ", ".join(equipment) if equipment else "Bodyweight only"

    intensity_map = {
        "Beginner":     "2–3 sets, moderate weight, 90s rest. Prioritise form.",
        "Intermediate": "3–4 sets, progressive overload, 60–75s rest.",
        "Advanced":     "4–5 sets, heavy compounds, 45–60s rest, supersets.",
    }
    intensity = intensity_map.get(fitness_level, "3 sets, 60s rest.")

    CHUNK_SIZE   = 3       # 3 days per call — well within Groq token limits
    chunks       = []
    start_day    = 1
    total_chunks = max(1, (total_days + CHUNK_SIZE - 1) // CHUNK_SIZE)

    while start_day <= total_days:
        end_day   = min(start_day + CHUNK_SIZE - 1, total_days)
        chunk_num = len(chunks) + 1

        if progress_callback:
            progress_callback(chunk_num, total_chunks, start_day - 1, total_days)

        context_note = ""
        if chunks:
            context_note = (
                f"(Continue from Day {start_day}. Do NOT repeat Days 1–{start_day-1}. "
                f"Vary muscle groups from previous days.)"
            )

        prompt = f"""Fitness trainer. Generate ONLY Days {start_day}-{end_day} of a {total_days}-day plan.
Client: {name}, {gender}, {height}cm/{weight}kg, BMI {bmi:.1f} ({bmi_cat}), Goal: {goal}, Level: {fitness_level}, Equipment: {eq_list}.
{context_note}

For EACH day use EXACTLY this format:
## Day N - [Muscle Group]
**Warm-Up** - Exercise: 2x10, Exercise: 2x10
**Main Workout**
- ExerciseName — 3x12 reps (rest 60s)
- ExerciseName — 3x12 reps (rest 60s)
- ExerciseName — 3x12 reps (rest 60s)
- ExerciseName — 3x12 reps (rest 60s)
- ExerciseName — 3x12 reps (rest 60s)
**Cool-Down** - Stretch1, Stretch2

Rules: vary muscle groups each day, {fitness_level} appropriate, use {eq_list} only, intensity: {intensity}.
{"End with 1 motivational sentence for "+name+"." if end_day == total_days else "No motivational text — more days follow."}
Output ONLY Days {start_day}-{end_day}. No preamble. No extra text.
"""

        # Retry up to 3 times on transient errors
        for attempt in range(3):
            try:
                chunk_text = query_model(prompt, max_tokens=1500)
                break
            except ValueError as e:
                # Rate limit — wait and retry
                err = str(e)
                if "rate limit" in err.lower() and attempt < 2:
                    wait = 65  # wait just over 1 minute for rate limit window to reset
                    if progress_callback:
                        progress_callback(chunk_num, total_chunks, start_day - 1, total_days,
                                          status=f"Rate limit — waiting {wait}s…")
                    time.sleep(wait)
                    continue
                raise
            except Exception as e:
                err = str(e)
                if attempt < 2:
                    time.sleep(5 * (attempt + 1))
                    continue
                raise

        chunks.append(chunk_text.strip())
        start_day = end_day + 1

        # Small pause between chunks to respect rate limits on free tier
        if start_day <= total_days:
            time.sleep(1)

    if progress_callback:
        progress_callback(total_chunks, total_chunks, total_days, total_days)

    full_plan = "\n\n".join(chunks)
    return full_plan, bmi, bmi_cat