"""Create: rephrase good source documents into new training documents, or write new ones where sources are thin.

Rephrase styles are the five from Sutra (weights 30/25/20/15/10, after BeyondWeb), plus math and code styles.
A fidelity check flags rephrasings that introduce numbers or names the source does not contain, which matters when the
training data must respect a knowledge cutoff.
"""
from __future__ import annotations

import random
import re

from .llm import Endpoint, EndpointError

NO_PREAMBLE = ("Output ONLY the rewritten document. No preamble, no meta-commentary, no phrases like \"Here is\", "
               "\"Let me\", \"Sure\" or \"Okay\". Do not add facts, numbers, names or dates that are not in the original.")

STYLES = {
    "conversational": ("Teacher-student dialogue", 30, """Rewrite the content as a natural conversation between a teacher and a student.
The student asks curious, natural questions; the teacher explains clearly with concrete examples taken from the content.
Keep all key information and the same technical accuracy. Keep a similar length (within 20%).
Use the labels **Teacher**: and **Student**: and start directly with "**Teacher**:"."""),
    "tutorial": ("Step-by-step tutorial", 25, """Rewrite the content as a clear step-by-step tutorial.
Start with "## Tutorial: <topic>", then an overview, prerequisites, numbered steps with what to do and how to check it,
a short practice exercise and key takeaways. Keep all key information accurate. Keep a similar length (within 20%)."""),
    "qa": ("Questions and answers", 20, """Rewrite the content as 4 to 7 questions and thorough answers that together cover every key point,
moving from basic (what) to deeper (why, how) to applied. Start with "## Understanding <topic>: Questions and Answers".
Use **Q:** and **A:** labels. Keep a similar total length (within 20%)."""),
    "dense_reference": ("Dense reference notes", 15, """Rewrite the content as compact reference notes at 70 to 85 percent of the original length.
Remove filler and repetition. Use headings and bullet points: core definition, key properties, relationships,
essential examples, formulas or procedures if any, pitfalls."""),
    "pedagogical": ("Textbook section", 10, """Rewrite the content as a well-structured textbook section: learning objectives, the connection to what
the reader already knows, the core explanation, one worked example with step-by-step reasoning, a common question
with its answer, and a quick self-check question. Keep a similar length."""),
    "worked_examples": ("Worked problems (math)", 0, """Rewrite the content as 2 to 4 worked problems that use exactly the ideas in it. For each: state the problem,
solve it step by step with the reasoning for each step, and check the answer. Use LaTeX for mathematics."""),
    "code_walkthrough": ("Code walkthrough (code)", 0, """Rewrite the content as an explained code walkthrough: what the code does, then the code in a fenced block
(keep it correct and runnable; you may tidy names and add comments), then an explanation of the key lines and one
example of using it."""),
}

KIND_STYLES = {"math": ["worked_examples", "pedagogical", "qa", "tutorial"],
               "code": ["code_walkthrough", "tutorial", "qa", "dense_reference"]}


def styles_for(kind: str) -> list[str]:
    return KIND_STYLES.get(kind, ["conversational", "tutorial", "qa", "dense_reference", "pedagogical"])


def pick_style(kind: str, rng: random.Random, allowed: list[str] | None = None) -> str:
    opts = [s for s in (allowed or styles_for(kind)) if s in STYLES]
    weights = [max(STYLES[s][1], 10) for s in opts]
    return rng.choices(opts, weights=weights, k=1)[0]


def chunks(text: str, max_chars: int = 6000) -> list[str]:
    """Split a long document at paragraph breaks so each rewrite sees a piece it can keep at a similar length."""
    if len(text) <= max_chars:
        return [text]
    out, cur = [], ""
    for para in re.split(r"\n\s*\n", text):
        if cur and len(cur) + len(para) + 2 > max_chars:
            out.append(cur)
            cur = ""
        cur = (cur + "\n\n" + para) if cur else para
        while len(cur) > max_chars:  # one giant paragraph
            out.append(cur[:max_chars])
            cur = cur[max_chars:]
    if cur.strip():
        out.append(cur)
    return [c for c in out if len(c.split()) >= 80] or [text[:max_chars]]


def rephrase_prompt(text: str, style: str, language: str = "") -> str:
    lang = f"\nWrite it in {language}." if language and language.lower() != "same" else ""
    return f"""{STYLES[style][2]}{lang}
{NO_PREAMBLE}

Original content:
<<<
{text[:8000]}
>>>"""


LEAD = re.compile(r"^(here('s| is)|sure|okay|certainly|let me|i('ll| will))\b", re.I)
NUM = re.compile(r"\b\d[\d,.]*\b")
# a capitalized word that follows a lowercase word mid-sentence: names, places, products. Words at the start of a
# sentence, a line or a markdown label are skipped, so "Teacher:" or "Great question" never count.
HEADINGS = re.compile(r"^#+ .*$|\*\*[^*\n]{1,80}\*\*", re.M)
NAME = re.compile(r"\b[a-z][a-z'-]*[,;]? ([A-Z][a-zA-Z]{2,}(?: [A-Z][a-zA-Z]{2,})*)")


def fidelity(src: str, out: str, template: str = "") -> dict:
    """Share of numbers and mid-sentence names in the rewrite that never appear in the source. A rough proxy for
    invented facts; the threshold is a choice, not a law."""
    s_low = (src + " " + template).lower()
    body = HEADINGS.sub(" ", out)  # headings and bold labels the style asks for are not new facts
    nums = {n.rstrip(".,") for n in NUM.findall(out)}
    names = set(NAME.findall(body))
    new_nums = [n for n in nums if n and n not in src]
    new_names = [n for n in names if n.lower() not in s_low]
    total = len(nums) + len(names)
    share = (len(new_nums) + len(new_names)) / total if total else 0.0
    return {"new_numbers": sorted(new_nums)[:8], "new_names": sorted(new_names)[:8], "new_share": round(share, 3),
            "checked": total}


def rephrase(ep: Endpoint, text: str, style: str, language: str = "", max_new_share: float = 0.35) -> dict:
    """One rewrite. Returns {"text", "style", "fidelity", "ok", "why"}."""
    def ok(s):
        return len(s.split()) >= 60 and not LEAD.search(s.strip()[:40])
    try:
        out = ep.chat(rephrase_prompt(text, style, language), max_tokens=3000, temperature=0.7, retries=3, validate=ok)
    except EndpointError as e:
        return {"text": "", "style": style, "ok": False, "why": f"model call failed: {e}"[:200]}
    fid = fidelity(text, out, STYLES[style][2])
    new = len(fid["new_numbers"]) + len(fid["new_names"])
    # invented numbers are the clearest sign of made-up content; one stray number or name is not a pattern
    good = len(fid["new_numbers"]) < 3 and (new < 2 or fid["new_share"] <= max_new_share)
    return {"text": out, "style": style, "fidelity": fid, "ok": good, "why": "" if good else "adds facts the source does not contain"}


# ---------------------------------------------------------------- new documents where sources are thin

DOC_TYPES = {  # a subset of Sutra's content types that reads well as pretraining text
    "concept_introduction": "an introduction to the concept: definition, intuition, why it matters, how it works, two examples, a common misconception, and how it connects to related ideas",
    "reasoning_demonstration": "a worked demonstration of reasoning about a concrete problem in this area, thinking step by step and checking the result",
    "worked_examples": "two or three worked examples of increasing difficulty, each solved step by step",
    "common_misconceptions": "an explanation of the most common misconceptions about the concept, why people hold them, and the correct understanding",
    "problem_set": "a short problem set of four problems with full solutions",
    "case_study": "a realistic case study that applies the concept, with the decisions and their reasons",
    "technical_documentation": "clear technical documentation of the concept or tool, with a usage example",
    "historical_context": "the history of how the idea developed, who contributed what, and why it mattered",
}


def invent_prompt(topic: str, doc_type: str, level: int, request: str, language: str = "English") -> str:
    return f"""You are writing high-quality pretraining text for a language model. Collection: "{request}".
Write {DOC_TYPES[doc_type]}.
Concept: {topic}
Difficulty: {level} on a 1 (introductory) to 10 (expert) scale. Language: {language}.
Write 350 to 700 words of dense, accurate prose with markdown headings where useful. Every sentence should teach
something. Do not invent statistics, quotes or citations. {NO_PREAMBLE.split('. Do not add')[0]}."""


def invent(ep: Endpoint, topic: str, doc_type: str, level: int, request: str, language: str = "English") -> dict:
    def ok(s):
        return len(s.split()) >= 150 and not LEAD.search(s.strip()[:40])
    try:
        out = ep.chat(invent_prompt(topic, doc_type, level, request, language), max_tokens=2500, temperature=0.8, retries=3, validate=ok)
    except EndpointError as e:
        return {"text": "", "ok": False, "why": f"model call failed: {e}"[:200]}
    return {"text": out, "ok": True, "doc_type": doc_type, "topic": topic, "level": level}