"""
groq_worker.py — Pipeline LLM Groq (free tier) pour désambiguïser les locations
que Geoapify ne sait pas géocoder après MAX_GEO_TRIES tentatives.

Flux :
  1. Sélectionner band_locations WHERE geocode_status='llm_needed'
  2. Vérifier llm_cache (sha256 de model+country+location_raw)
  3. Sinon appeler Groq (JSON mode garanti) → extraire city + iso2
  4. Si ville trouvée → geocode_query=clean_query, status='queued', tries_geo=0
  5. Si aucune ville (is_null) → incrémenter tries_llm, backoff exponentiel
  6. Après MAX_LLM_TRIES → status='manual'

Rate limits Groq free tier :
  llama-3.3-70b-versatile : 30 req/min, 1 000 req/jour
  llama-3.1-8b-instant    : 30 req/min, 14 400 req/jour (fallback)
"""

import hashlib
import json
import os
import re
import sys
import time
from datetime import UTC, datetime

import psycopg2
import requests

# sys.path[0] = src/ quand lancé comme "python src/groq_worker.py"
from parser import COUNTRY_NAMES

GROQ_API_KEY  = os.environ.get("GROQ_API_KEY", "").strip()
GROQ_BASE     = "https://api.groq.com/openai/v1/chat/completions"
MAX_LLM_TRIES = int(os.environ.get("GROQ_MAX_TRIES", "5"))
CALL_DELAY    = float(os.environ.get("GROQ_CALL_DELAY", "2.5"))

MODELS = [
    # (model_id, req_per_min, req_per_day)
    ("llama-3.3-70b-versatile", 30, 1000),
    ("llama-3.1-8b-instant",    30, 14400),
]

# Tarifs Groq approximatifs par modèle : (usd/token_in, usd/token_out)
PRICING = {
    "llama-3.3-70b-versatile": (0.59e-6, 0.79e-6),
    "llama-3.1-8b-instant":    (0.05e-6, 0.08e-6),
}


def _cost(model: str, tok_in: int, tok_out: int) -> float:
    pin, pout = PRICING.get(model, (0.59e-6, 0.79e-6))
    return tok_in * pin + tok_out * pout

_SYSTEM = (
    "You are a precise location data extraction assistant. "
    "Your only job is to extract a city name and ISO 3166-1 alpha-2 country code "
    "from a location description. Always respond with valid JSON only."
)

_USER_TMPL = """\
Extract the primary city (or town/village) and ISO 3166-1 alpha-2 country code \
from this location.

Band's registered country: {country}
Location text: "{location}"

Rules:
- Use the band's country as context if the location is ambiguous
- Return the most specific city/town name you can extract
- Strip all administrative divisions: keep just the city name
- If there are multiple cities separated by "/" or "and", return the first one
- If no real place can be extracted, set city to null

Respond ONLY with this JSON (no other text):
{{"city": "city name or null", "country": "ISO2 code", "confidence": 0.0}}"""


def _input_hash(model: str, location_raw: str, country: str) -> str:
    key = f"{model}|{(country or '').upper()}|{location_raw.strip().lower()}"
    return hashlib.sha256(key.encode()).hexdigest()


def _call_groq(model: str, location_raw: str, country: str) -> tuple[str, int, int]:
    prompt = _USER_TMPL.format(
        country=country or "unknown (European metal band)",
        location=location_raw,
    )
    r = requests.post(
        GROQ_BASE,
        headers={
            "Authorization": f"Bearer {GROQ_API_KEY}",
            "Content-Type":  "application/json",
        },
        json={
            "model":           model,
            "messages":        [
                {"role": "system", "content": _SYSTEM},
                {"role": "user",   "content": prompt},
            ],
            "max_tokens":      80,
            "temperature":     0.0,
            "response_format": {"type": "json_object"},
        },
        timeout=30,
    )
    if r.status_code == 429:
        raise RuntimeError(f"rate_limit:{model}")
    r.raise_for_status()
    data  = r.json()
    text  = data["choices"][0]["message"]["content"]
    usage = data.get("usage", {})
    return text, usage.get("prompt_tokens", 0), usage.get("completion_tokens", 0)


def _parse_json(text: str) -> dict | None:
    try:
        return json.loads(text.strip())
    except json.JSONDecodeError:
        # Repli quand le modèle enrobe son JSON de texte. `[^}]+` s'arrêtait au
        # PREMIER `}` : dès que la réponse contenait un objet imbriqué, la
        # capture était tronquée et invalide, le lieu comptait pour un échec et
        # consommait un essai LLM pour rien. Le re.DOTALL était de surcroît
        # inopérant, `[^}]` matchant déjà les retours à la ligne.
        # `.*` glouton va du premier `{` au dernier `}` : correct pour un objet
        # unique noyé dans de la prose, ce qui est le seul cas observé.
        m = re.search(r'\{.*\}', text, re.DOTALL)
        if m:
            try:
                return json.loads(m.group())
            except json.JSONDecodeError:
                pass
    return None


def _apply_clean_query(cur, loc_id: int, city: str, iso2: str | None, tries_llm: int):
    country_name = COUNTRY_NAMES.get((iso2 or '').upper(), iso2 or '')
    clean_query  = f"{city}, {country_name}" if country_name else city
    cur.execute(
        """
        UPDATE band_locations
        SET geocode_status='queued',
            geocode_query=%s,
            geocode_tries_geo=0,
            geocode_tries_llm=%s,
            geocode_next_at=now(),
            geocode_error=NULL,
            updated_at=now()
        WHERE id=%s
        """,
        (clean_query, tries_llm + 1, loc_id),
    )


def main():
    if not GROQ_API_KEY:
        print("[groq] GROQ_API_KEY non configuré — arrêt")
        sys.exit(0)

    dsn  = os.environ["DATABASE_URL"]
    conn = psycopg2.connect(dsn)
    conn.autocommit = True

    day_count = {m[0]: 0 for m in MODELS}
    min_count = {m[0]: 0 for m in MODELS}
    min_start = time.monotonic()
    day_start = datetime.now(UTC).date()

    with conn.cursor() as cur:
        while True:
            # Reset compteurs si nouvelle minute / nouveau jour
            if time.monotonic() - min_start >= 60:
                min_count = {m[0]: 0 for m in MODELS}
                min_start = time.monotonic()
            today = datetime.now(UTC).date()
            if today != day_start:
                day_count = {m[0]: 0 for m in MODELS}
                day_start = today

            cur.execute(
                """
                SELECT bl.id, bl.ma_id, bl.location_raw, bl.geocode_tries_llm,
                       b.country
                FROM band_locations bl
                JOIN bands b ON b.ma_id = bl.ma_id
                WHERE bl.geocode_status = 'llm_needed'
                  AND bl.geocode_next_at <= now()
                ORDER BY bl.geocode_tries_llm ASC, bl.id ASC
                LIMIT 1
                FOR UPDATE OF bl SKIP LOCKED
                """
            )
            row = cur.fetchone()
            if not row:
                print("[groq] rien à traiter, attente 120s")
                time.sleep(120)
                continue

            loc_id, ma_id, location_raw, tries_llm, country = row

            # Choisir un modèle disponible
            chosen_model = None
            for model, rpm, rpd in MODELS:
                if min_count[model] < rpm and day_count[model] < rpd:
                    chosen_model = model
                    break

            if not chosen_model:
                wait = max(5, 60 - (time.monotonic() - min_start))
                print(f"[groq] quota atteint, attente {wait:.0f}s")
                cur.execute(
                    "UPDATE band_locations SET geocode_status='llm_needed', updated_at=now() WHERE id=%s",
                    (loc_id,),
                )
                time.sleep(wait)
                continue

            # Vérifier llm_cache
            cache_hash = _input_hash(chosen_model, location_raw, country or '')
            cur.execute(
                "SELECT parsed_city, parsed_country, is_null FROM llm_cache WHERE input_hash=%s",
                (cache_hash,),
            )
            cached = cur.fetchone()

            if cached:
                parsed_city, parsed_country, is_null = cached
                print(f"[groq] id={loc_id} CACHE HIT city={parsed_city}")
                if is_null or not parsed_city:
                    _handle_null_city(cur, loc_id, tries_llm)
                else:
                    _apply_clean_query(cur, loc_id, parsed_city, parsed_country, tries_llm)
                continue

            # Appel API Groq
            try:
                response_text, tok_in, tok_out = _call_groq(
                    chosen_model, location_raw, country or ''
                )
                min_count[chosen_model] += 1
                day_count[chosen_model] += 1

                parsed      = _parse_json(response_text)
                city        = (parsed or {}).get("city")
                iso2        = (parsed or {}).get("country") or country
                is_null     = not city
                cost_usd    = _cost(chosen_model, tok_in, tok_out)

                prompt_text = _USER_TMPL.format(
                    country=country or "unknown (European metal band)",
                    location=location_raw,
                )
                cur.execute(
                    """
                    INSERT INTO llm_cache
                      (input_hash, model, prompt, response,
                       parsed_city, parsed_country, is_null,
                       tokens_in, tokens_out, cost_usd,
                       ma_id, location_raw, country)
                    VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                    ON CONFLICT (input_hash) DO NOTHING
                    """,
                    (cache_hash, chosen_model, prompt_text, response_text,
                     city, iso2, is_null, tok_in, tok_out, cost_usd,
                     ma_id, location_raw, country),
                )

                if is_null:
                    print(f"[groq] id={loc_id} city=null (tries_llm={tries_llm + 1})")
                    _handle_null_city(cur, loc_id, tries_llm)
                else:
                    _apply_clean_query(cur, loc_id, city, iso2, tries_llm)
                    print(f"[groq] id={loc_id} → queued city='{city}' ({iso2})")

            except RuntimeError as exc:
                if "rate_limit" in str(exc):
                    hit_model = str(exc).split(":", 1)[1] if ":" in str(exc) else chosen_model
                    min_count[hit_model] = 9999
                    print(f"[groq] rate limit {hit_model}")
                    cur.execute(
                        "UPDATE band_locations SET geocode_status='llm_needed', updated_at=now() WHERE id=%s",
                        (loc_id,),
                    )
                    continue
                # Autre erreur réseau/API → backoff
                new_tries = tries_llm + 1
                backoff   = 30 * new_tries
                cur.execute(
                    """
                    UPDATE band_locations
                    SET geocode_tries_llm=%s,
                        geocode_error=%s,
                        geocode_next_at=now() + (%s || ' minutes')::interval,
                        updated_at=now()
                    WHERE id=%s
                    """,
                    (new_tries, str(exc)[:300], backoff, loc_id),
                )
                print(f"[groq] id={loc_id} error: {exc}")

            time.sleep(CALL_DELAY)

    conn.close()


def _handle_null_city(cur, loc_id: int, tries_llm: int):
    new_tries = tries_llm + 1
    if new_tries >= MAX_LLM_TRIES:
        cur.execute(
            """
            UPDATE band_locations
            SET geocode_status='manual',
                geocode_error='llm: city=null après max tries',
                geocode_tries_llm=%s,
                updated_at=now()
            WHERE id=%s
            """,
            (new_tries, loc_id),
        )
        print(f"[groq] id={loc_id} → manual (exhausted)")
    else:
        backoff = 60 * new_tries
        cur.execute(
            """
            UPDATE band_locations
            SET geocode_tries_llm=%s,
                geocode_next_at=now() + (%s || ' minutes')::interval,
                updated_at=now()
            WHERE id=%s
            """,
            (new_tries, backoff, loc_id),
        )


if __name__ == "__main__":
    main()