"""DataInventor engine: one sentence in, a training dataset out, on any OpenAI-compatible endpoint.

Pipeline: blueprint (taxonomy + personas + question shapes) -> sampled combinations ->
question -> answer (or two answers + a ranking for preference pairs) -> filters -> near-duplicate rejection.
"""
from __future__ import annotations

import json
import random
import re
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import Optional

from datainventor.llm import HF_BASE, HF_MODEL, Endpoint, EndpointError  # noqa: F401  (main imports them from here)


# ---------------------------------------------------------------- blueprint

def _s(x) -> str:
    """Models sometimes return {"persona": "..."} instead of a bare string."""
    if isinstance(x, dict):
        for k in ("name", "persona", "role", "format", "shape", "topic", "description", "title"):
            if isinstance(x.get(k), str):
                return x[k].strip()
        return " ".join(str(v) for v in x.values() if isinstance(v, str)).strip()
    return str(x).strip()


def normalize_blueprint(bp: dict) -> dict:
    out = {"topics": [], "personas": [_s(p) for p in bp.get("personas", []) if _s(p)],
           "formats": [_s(f) for f in bp.get("formats", []) if _s(f)],
           "difficulty": [_s(d) for d in bp.get("difficulty", []) if _s(d)] or ["basic", "intermediate", "advanced"],
           "assistant_role": _s(bp.get("assistant_role") or "an excellent, accurate AI assistant")[:300]}
    for a in bp.get("topics", []):
        subs = []
        for sub in a.get("subtopics", []):
            leaves = [_s(l) for l in sub.get("leaf_topics", []) if _s(l)]
            if leaves:
                subs.append({"subtopic": _s(sub.get("subtopic", "")), "leaf_topics": leaves})
        if subs:
            out["topics"].append({"area": _s(a.get("area", "")), "subtopics": subs})
    if len(out["topics"]) < 1 or len(out["personas"]) < 2 or len(out["formats"]) < 2:
        raise ValueError("the blueprint is missing topics, personas or formats")
    return out


def _json_from(text: str):
    m = re.search(r"\{.*\}", text, re.S)
    if not m:
        raise ValueError("no JSON object in the reply")
    return json.loads(m.group(0))


def make_blueprint(ep: Endpoint, request: str, domains: list[str], subdomains: list[str], languages: list[str],
                   constraints: str, rows: int) -> dict:
    scope = ""
    if domains:
        scope += f"Domains: {', '.join(domains)}.\n"
    if subdomains:
        scope += f"Subdomains: {', '.join(subdomains)}.\n"
    if constraints.strip():
        scope += f"Output requirements: {constraints.strip()}\n"
    areas = 6 if rows <= 500 else 8
    prompt = f"""You are planning a training dataset. Dataset request: "{request}"
{scope}
Return ONLY JSON with these keys:
"topics": a list of {areas} objects {{"area": str, "subtopics": [4 objects {{"subtopic": str, "leaf_topics": [4 specific, concrete strings]}}]}}. Cover the whole request evenly.
"personas": 10 different kinds of people who would write the user side of this dataset (for a customer-support dataset, 10 different kinds of customer; for a coding dataset, 10 kinds of developer).
"assistant_role": one sentence describing who answers in this dataset and in what voice (for example "a friendly support agent at a small online bookshop who gives concrete next steps").
"formats": 6 different shapes the instruction could take (for example a short direct ask, a detailed brief, a request with a concrete example, a request with constraints, a follow-up, a comparison).
"difficulty": ["basic","intermediate","advanced"].
No commentary, no code fences."""
    last = None
    for _ in range(4):
        try:
            return normalize_blueprint(_json_from(ep.chat(prompt, max_tokens=5000, temperature=0.7)))
        except EndpointError:
            raise
        except Exception as e:
            last = e
    raise EndpointError(f"The model did not return a usable blueprint: {last}")


# ---------------------------------------------------------------- generation

LEAK = re.compile(r"the user (asks|wants)|write one|return only|^okay|^sure,|^here is|^here's|^as an ai", re.I)
END = re.compile(r"[.?!)\"”’'।:。！？]$")
MULTI = re.compile(r"\b(generate|draft|write|compile|create|give me|produce)\s+(\d+|several|multiple|a set of|a list of)\b", re.I)
PLACEHOLDER = re.compile(r"\[[A-Za-z'’ ]{0,30}\b(name|number|email|address|phone|date|id)\]", re.I)
CALLBACK = re.compile(r"\b(like|as) you (said|mentioned|suggested|told me)\b|\b(last time|earlier) (you|we)\b|\bthanks again\b", re.I)
TABLE = re.compile(r"^\s*\|.*\|\s*$\n^\s*\|?\s*:?-{3,}", re.M)
VERDICT = re.compile(r"\W*\(?([AB])\b", re.I)
LENGTHS = ["one blunt sentence", "two sentences with one specific detail",
           "three or four sentences with concrete specifics", "a short paragraph that sounds a little rushed"]


def degenerate(t: str) -> bool:
    w = t.lower().split()
    return bool(re.search(r"(.)\1{7,}", t)) or (len(w) >= 8 and len(set(w)) / len(w) < 0.5)


def grams(t: str, n: int = 3) -> set:
    w = t.lower().split()
    return set(zip(*[w[i:] for i in range(n)])) if len(w) >= n else {tuple(w)}


@dataclass
class Job:
    id: str
    config: dict
    state: str = "queued"  # queued | planning | generating | done | failed | stopped
    message: str = ""
    blueprint: Optional[dict] = None
    rows: list = field(default_factory=list)
    target: int = 0
    attempts: int = 0
    rejected: dict = field(default_factory=dict)
    started: float = field(default_factory=time.time)
    finished: float = 0.0
    stop: threading.Event = field(default_factory=threading.Event)
    lock: threading.Lock = field(default_factory=threading.Lock)

    def reject(self, why: str):
        with self.lock:
            self.rejected[why] = self.rejected.get(why, 0) + 1


def run_job(job: Job, ep: Endpoint):
    cfg = job.config
    rows_n = int(cfg["rows"])
    job.target = rows_n
    kind = cfg.get("format", "instruction")
    request = cfg["request"]
    languages = cfg.get("languages") or ["English"]
    constraints = cfg.get("constraints", "")
    try:
        job.state, job.message = "planning", "Drafting the blueprint"
        bp = normalize_blueprint(cfg["blueprint"]) if cfg.get("blueprint") else make_blueprint(
            ep, request, cfg.get("domains", []), cfg.get("subdomains", []), languages, constraints, rows_n)
        job.blueprint = bp
        leaves = [(a["area"], s["subtopic"], l) for a in bp["topics"] for s in a["subtopics"] for l in s["leaf_topics"]]
        rng = random.Random(cfg.get("seed", 7))
        combos, seen = [], set()
        want = int(rows_n * 1.8) + 20
        guard = 0
        while len(combos) < want and guard < want * 30:
            guard += 1
            c = (rng.choice(leaves), rng.choice(bp["personas"]), rng.choice(bp["formats"]), rng.choice(bp["difficulty"]),
                 rng.choice(LENGTHS), rng.choice(languages))
            if c not in seen:
                seen.add(c)
                combos.append(c)

        scope = f'The dataset being built: "{request}"'
        if constraints.strip():
            scope += f"\nEvery example must satisfy: {constraints.strip()}"

        def ok_q(t):
            t = t.strip().strip('"')
            return (15 <= len(t) <= 1500 and END.search(t) and not degenerate(t) and not LEAK.search(t[:150])
                    and not MULTI.search(t[:120]) and not PLACEHOLDER.search(t) and not CALLBACK.search(t))

        def ok_a(t):
            return (len(t.split()) >= 12 and not degenerate(t[:600]) and not LEAK.search(t[:100])
                    and not TABLE.search(t) and not PLACEHOLDER.search(t))

        kept_grams: list[set] = []
        job.state, job.message = "generating", "Writing examples"

        def work(idx_combo):
            idx, ((area, sub, leaf), persona, fmt, diff, length, lang) = idx_combo
            if job.stop.is_set():
                return
            with job.lock:
                if len(job.rows) >= rows_n:
                    return
                job.attempts += 1
            try:
                role = bp.get("assistant_role", "an excellent AI assistant")
                q = ep.chat(f"""{scope}
Write ONE realistic message that {persona} would send to {role}, for this dataset, about: {leaf} ({area}, {sub}).
Shape: {fmt}. Difficulty: {diff}. Length: {length}. Language: {lang}.
It should sound like a real person typing, specific and a little imperfect, not a textbook prompt. It must be ONE question or ONE task whose answer fits in a single reply: never ask for a list of many items, a set of Q&A pairs, or a dataset. Do not name the shape or the persona.
It is the first message in the conversation: do not refer to earlier chats or to anything said before. Do not use placeholders such as [Name] or [order number].
Return only the message.""", max_tokens=700, temperature=1.0, validate=ok_q).strip().strip('"')
                g = grams(q)
                with job.lock:
                    if any(len(g & k) / max(1, len(g | k)) > 0.5 for k in kept_grams):
                        job.reject("near-duplicate question")
                        return
                answer_prompt = f"""You are {role}. Reply to the message below. The reply is a training example for this dataset: "{request}".
Language: {lang}. Be accurate, specific and complete. Start directly with the content, with no preamble about who you are. Do not invent statistics or facts you cannot know.
Write plain prose: short paragraphs, or a simple numbered list for steps. No tables, headings or emoji.
Never invent order numbers, dates, prices, phone numbers, emails, URLs, product or brand names, menu paths or policies. If a detail like that is needed, ask for it or speak generally. Do not use placeholders such as [Your Name].
For physical techniques and safety, give only the standard, widely taught practice, and make sure every step agrees with the others.
{('Requirements: ' + constraints.strip()) if constraints.strip() else ''}

Message: {q}"""
                a1 = ep.chat(answer_prompt, max_tokens=2500, temperature=0.7, validate=ok_a)
                row = {"area": area, "subtopic": sub, "topic": leaf, "persona": persona, "shape": fmt,
                       "difficulty": diff, "language": lang, "prompt": q}
                if kind == "preference":
                    a2 = ep.chat(answer_prompt + "\n\nWrite a different answer from the obvious one, in a different structure and with different emphasis.",
                                 max_tokens=2500, temperature=1.0, validate=ok_a)
                    def judge(first: str, second: str) -> str:
                        v = ep.chat(f"""Which answer is better for the message? Judge accuracy first: an answer that invents details (order numbers, dates, contact details, policies) or gives unsafe advice is worse. Then completeness and clarity.
Message: {q}

Answer A:
{first[:3000]}

Answer B:
{second[:3000]}

Reply with only the single letter A or B.""", max_tokens=16, temperature=0, allow_cut=True,
                                    validate=lambda s: bool(VERDICT.match(s)))
                        return VERDICT.match(v).group(1).upper()
                    # judge both orders so the first-shown answer gets no edge; keep the row only if they agree
                    v1, v2 = judge(a1, a2), judge(a2, a1)
                    if (v1 == "A") != (v2 == "B"):
                        job.reject("judge disagreed")
                        return
                    row["chosen"], row["rejected"] = (a1, a2) if v1 == "A" else (a2, a1)
                else:
                    row["completion"] = a1
                with job.lock:
                    if len(job.rows) >= rows_n:
                        return
                    if any(len(g & k) / max(1, len(g | k)) > 0.5 for k in kept_grams):
                        job.reject("near-duplicate question")
                        return
                    row["id"] = len(job.rows) + 1
                    job.rows.append(row)
                    kept_grams.append(g)
                    job.message = f"{len(job.rows)} of {rows_n} written"
                    if len(job.rows) >= rows_n:
                        job.state, job.message, job.finished = "done", f"{len(job.rows)} rows ready", time.time()
                        job.stop.set()
            except EndpointError as e:
                if "refused the key" in str(e) or "not found" in str(e):
                    job.state, job.message = "failed", str(e)
                    job.stop.set()
                else:
                    job.reject("model call failed")
            except Exception:
                job.reject("model call failed")

        with ThreadPoolExecutor(max_workers=max(1, min(ep.workers, 16))) as pool:
            list(pool.map(work, enumerate(combos)))
        if job.state in ("failed", "done"):
            return
        if job.stop.is_set():
            job.state, job.message = "stopped", f"Stopped with {len(job.rows)} rows"
        elif len(job.rows) < rows_n:
            job.state = "done"
            job.message = f"Finished with {len(job.rows)} of {rows_n} rows. The model rejected too many attempts to fill the rest."
        else:
            job.state, job.message = "done", f"{len(job.rows)} rows ready"
    except EndpointError as e:
        job.state, job.message = "failed", str(e)
    except Exception as e:  # never leave a job hanging
        job.state, job.message = "failed", f"Unexpected error: {e}"
    finally:
        job.finished = time.time()


# ---------------------------------------------------------------- stats

def stats(rows: list) -> dict:
    if not rows:
        return {}
    text = "\n".join(r["prompt"] for r in rows).lower().split()
    tri = list(zip(text, text[1:], text[2:]))
    area_counts: dict = {}
    for r in rows:
        area_counts[r["area"]] = area_counts.get(r["area"], 0) + 1
    comp = [len((r.get("completion") or r.get("chosen") or "").split()) for r in rows]
    return {"rows": len(rows), "distinct3_prompts": round(len(set(tri)) / max(1, len(tri)), 3),
            "avg_prompt_words": round(sum(len(r["prompt"].split()) for r in rows) / len(rows), 1),
            "avg_answer_words": round(sum(comp) / len(rows), 1), "areas": area_counts}