#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Extraction des specs de propriété via GPT-4o-mini, à partir des descriptions.
Parallèle, résumable (cache specs.jsonl), avec retry/backoff. Aucune dépendance.
"""
import json, os, time, threading, urllib.request, urllib.error
from concurrent.futures import ThreadPoolExecutor, as_completed

SRC = "data/processed/maisons_avec_prix.json"
OUT = "data/processed/specs.jsonl"
WORKERS = int(os.environ.get("WORKERS", "24"))
NSHARDS = int(os.environ.get("NSHARDS", "1"))
SHARD = int(os.environ.get("SHARD", "0"))
MODEL = "gpt-4o-mini"

# clé
for line in open(".env", encoding="utf-8"):
    if line.startswith("OPENAI_API_KEY="):
        os.environ["OPENAI_API_KEY"] = line.split("=", 1)[1].strip()
KEY = os.environ["OPENAI_API_KEY"]

SCHEMA = {
    "type": "object", "additionalProperties": False,
    "properties": {
        "condition": {"type": "string", "enum": ["a_renover","original","bon","renove","refait_neuf","neuf","inconnu"]},
        "turnkey": {"type": "boolean"},
        "recently_renovated": {"type": "boolean"},
        "flooring": {"type": "array", "items": {"type": "string", "enum": ["bois_franc","ceramique","plancher_flottant","tapis","vinyle","beton","autre"]}},
        "countertop": {"type": "string", "enum": ["granit","quartz","marbre","stratifie","bois","autre","inconnu"]},
        "heating": {"type": "array", "items": {"type": "string", "enum": ["electrique","plinthes","thermopompe","bois","gaz","mazout","radiant","geothermie","autre"]}},
        "garage": {"type": "boolean"},
        "pool": {"type": "boolean"},
        "waterfront": {"type": "boolean"},
        "view": {"type": "string", "enum": ["eau","montagne","ville","aucune","inconnu"]},
        "finished_basement": {"type": "boolean"},
        "fireplace": {"type": "boolean"},
        "central_ac": {"type": "boolean"},
        "new_construction": {"type": "boolean"},
        "income_property": {"type": "boolean"},
        "renovated_kitchen": {"type": "boolean"},
        "renovated_bathroom": {"type": "boolean"},
        "large_lot": {"type": "boolean"},
        "quiet_area": {"type": "boolean"},
        "highlights": {"type": "array", "items": {"type": "string"}}
    },
    "required": ["condition","turnkey","recently_renovated","flooring","countertop",
                 "heating","garage","pool","waterfront","view","finished_basement",
                 "fireplace","central_ac","new_construction","income_property",
                 "renovated_kitchen","renovated_bathroom","large_lot","quiet_area",
                 "highlights"]
}
SYS = ("Tu es un extracteur d'attributs immobiliers. À partir de la description "
       "d'une propriété (français ou anglais), renvoie les specs au format JSON "
       "strict. Base-toi UNIQUEMENT sur le texte ; si une info est absente, mets "
       "false / 'inconnu' / [] selon le champ. highlights = 3 à 6 mots-clés de vente.")

def call(desc, tries=5):
    body = {"model": MODEL, "temperature": 0,
            "messages": [{"role": "system", "content": SYS},
                         {"role": "user", "content": desc[:4000]}],
            "response_format": {"type": "json_schema", "json_schema":
                {"name": "specs", "strict": True, "schema": SCHEMA}}}
    data = json.dumps(body).encode()
    for t in range(tries):
        try:
            req = urllib.request.Request(
                "https://api.openai.com/v1/chat/completions", data=data,
                headers={"Authorization": f"Bearer {KEY}",
                         "Content-Type": "application/json"})
            with urllib.request.urlopen(req, timeout=90) as r:
                out = json.load(r)
            return json.loads(out["choices"][0]["message"]["content"])
        except urllib.error.HTTPError as e:
            if e.code in (429, 500, 502, 503, 504) and t < tries - 1:
                time.sleep(2 ** t + 0.5); continue
            raise
        except (urllib.error.URLError, TimeoutError):
            if t < tries - 1:
                time.sleep(2 ** t + 0.5); continue
            raise
    return None

# charge données + cache
data = json.load(open(SRC, encoding="utf-8"))
done = set()
if os.path.exists(OUT):
    for line in open(OUT, encoding="utf-8"):
        try:
            done.add(json.loads(line)["uid"])
        except Exception:
            pass

todo = []
for d in data:
    uid = f"{d.get('source')}:{d.get('listingId')}"
    desc = d.get("description")
    if uid in done or not desc:
        continue
    todo.append((uid, desc))

# sharding : partition disjointe et complète par index (déterministe)
if NSHARDS > 1:
    todo = [t for i, t in enumerate(todo) if i % NSHARDS == SHARD]

print(f"[shard {SHARD}/{NSHARDS}, workers={WORKERS}] Total={len(data)}  "
      f"déjà fait={len(done)}  à traiter (ce shard)={len(todo)}", flush=True)

lock = threading.Lock()
counter = {"n": 0, "err": 0}
fh = open(OUT, "a", encoding="utf-8")

def work(item):
    uid, desc = item
    try:
        specs = call(desc)
        rec = {"uid": uid, "specs": specs}
    except Exception as e:
        rec = {"uid": uid, "error": str(e)[:200]}
    with lock:
        fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
        fh.flush()
        counter["n"] += 1
        if "error" in rec:
            counter["err"] += 1
        if counter["n"] % 250 == 0:
            print(f"  {counter['n']}/{len(todo)}  (err={counter['err']})",
                  flush=True)

t0 = time.time()
with ThreadPoolExecutor(max_workers=WORKERS) as ex:
    list(ex.map(work, todo))
fh.close()
print(f"Terminé : {counter['n']} traités, {counter['err']} erreurs, "
      f"{time.time()-t0:.0f}s", flush=True)