import argparse
import json
import os
import re
import statistics
from collections import Counter
from urllib import request, error


FOREIGN_INSERTIONS = [
    "pero en espanol dice que no retrocedas",
    "y tambien dijo que Dios responde",
    "mas ele continuou em fe",
    "et il a dit de rester ferme",
]


def parse_distribute_debug_log(path):
    entries = []
    current = {"final": None, "en_bridge": None, "llm_paragraph": None}

    with open(path, "r", encoding="utf-8") as handle:
        for raw_line in handle:
            line = raw_line.rstrip("\n")
            if "[Final]:" in line:
                current["final"] = line.split("[Final]:", 1)[1].strip()
            elif "[EN-BRIDGE]:" in line:
                current["en_bridge"] = line.split("[EN-BRIDGE]:", 1)[1].strip()
            elif "[LLM-PARAGRAPH]:" in line:
                current["llm_paragraph"] = line.split("[LLM-PARAGRAPH]:", 1)[1].strip()
                entries.append(current)
                current = {"final": None, "en_bridge": None, "llm_paragraph": None}
            elif not line.strip():
                continue

    return entries


def normalize_paragraph(text):
    return (text or "").strip().strip('"').strip()


def reconstruct_samples(path):
    samples = []
    rows = parse_distribute_debug_log(path)
    previous_llm = ""
    pending_bridges = []
    recent_raw_windows = []

    for row in rows:
        if row.get("en_bridge"):
            pending_bridges.append(row["en_bridge"])
        if not row.get("llm_paragraph"):
            continue

        llm_output = normalize_paragraph(row["llm_paragraph"])
        new_raw = " ".join(pending_bridges).strip()
        prev_raw = recent_raw_windows[-1] if recent_raw_windows else ""

        samples.append(
            {
                "previous_refined": previous_llm,
                "previous_raw": prev_raw,
                "new_raw": new_raw,
                "observed_output": llm_output,
            }
        )

        previous_llm = llm_output
        recent_raw_windows.append(new_raw)
        pending_bridges = []

    return samples


def inject_multilingual_noise(text, sample_index, slot_name):
    if not text:
        return text
    inserts = [
        FOREIGN_INSERTIONS[(sample_index + 0) % len(FOREIGN_INSERTIONS)],
        FOREIGN_INSERTIONS[(sample_index + 1) % len(FOREIGN_INSERTIONS)],
    ]
    sentences = [part.strip() for part in re.split(r"(?<=[.!?])\s+", text.strip()) if part.strip()]
    if not sentences:
        return text
    if len(sentences) == 1:
        return f'{sentences[0]} {inserts[0]}.'
    head = sentences[0]
    tail = " ".join(sentences[1:])
    if slot_name == "previous_raw":
        return f"{head} {inserts[0]}. {tail}"
    return f"{head} {tail} {inserts[1]}."


def maybe_augment_samples(samples, mode):
    if mode != "multilingual_noise":
        return samples

    augmented = []
    for idx, sample in enumerate(samples):
        updated = dict(sample)
        updated["previous_raw"] = inject_multilingual_noise(sample["previous_raw"], idx, "previous_raw")
        updated["new_raw"] = inject_multilingual_noise(sample["new_raw"], idx, "new_raw")
        augmented.append(updated)
    return augmented


def build_variants(sample, variant_filter=None):
    prev_refined = sample["previous_refined"]
    prev_raw = sample["previous_raw"]
    new_raw = sample["new_raw"]

    variants = []

    system_a = """You are a live transcript editor.

Task:
Refine live transcription stream into clean, professional paragraphs.

Rules:
1. Integrate new segments into the flow of the previous working context.
2. Remove redundant repetitions and translator echoes.
3. Use a double newline (\\n\\n) to start a new paragraph when a topic changes or the current one is complete.
4. Output only in English.
5. Return only the refined, consolidated text with no explanation."""
    user_a = f"""Refine the following live transcription stream.

[PREVIOUS WORKING CONTEXT]
"{prev_refined}"

[NEW SOURCE TEXT]
"{new_raw}"
"""
    variants.append(("current_refined_plus_new", system_a, user_a))

    system_b = """You are a live transcript editor for noisy speech recognition output.

Goal:
Produce a clean English paragraph update grounded in the source text.

Grounding rules:
1. Treat PREVIOUS REFINED CONTEXT as style and continuity help, not as unquestionable truth.
2. Use PREVIOUS SOURCE TEXT and NEW SOURCE TEXT as the factual anchor.
3. Do not add details that are not supported by the source text.
4. Preserve uncertainty when the source is unclear instead of inventing a cleaner claim.
5. Remove obvious repetitions and translator echoes.
6. Start a new paragraph with a double newline only when there is a clear topic shift.
7. Return only the refined English text with no explanation."""
    user_b = f"""Update the transcript using the source text below.

[PREVIOUS REFINED CONTEXT]
"{prev_refined}"

[PREVIOUS SOURCE TEXT]
"{prev_raw}"

[NEW SOURCE TEXT]
"{new_raw}"
"""
    variants.append(("refined_plus_prev_raw_plus_new_raw", system_b, user_b))

    system_c = """You are a careful transcript editor.

Goal:
Produce clean English paragraphs directly from noisy source text.

Rules:
1. Use only the supplied source text as evidence.
2. Do not infer missing facts.
3. Keep wording conservative when the source is noisy.
4. Remove repeated fragments and obvious translation artifacts.
5. Use a double newline only for a clear paragraph break.
6. Return only the refined English text."""
    user_c = f"""Refine this transcript source.

[RECENT SOURCE TEXT]
"{prev_raw}"

[NEW SOURCE TEXT]
"{new_raw}"
"""
    variants.append(("prev_raw_plus_new_raw_only", system_c, user_c))

    system_d = """You are a live transcript editor for sermon audio translated into English.

Goal:
Create the most faithful readable English paragraph you can from imperfect source text.

Priority order:
1. Faithfulness to source text.
2. Remove duplicated phrases, stutters, and translation echoes.
3. Maintain continuity with the previous refined context only when it matches the source text.
4. Prefer slight awkwardness over hallucination.
5. If a phrase is unclear, keep it modestly literal instead of guessing.
6. Use a double newline only for a real topic change.
7. Return only the revised English transcript.

Hard constraints:
- Do not add people, events, or meanings absent from the source text.
- Do not turn rhetorical questions into factual claims unless the source clearly does that.
- Do not overwrite NEW SOURCE TEXT with PREVIOUS REFINED CONTEXT."""
    user_d = f"""Revise the transcript update.

[PREVIOUS REFINED CONTEXT]
"{prev_refined}"

[RECENT SOURCE WINDOW]
"{prev_raw}"

[CURRENT SOURCE WINDOW]
"{new_raw}"
"""
    variants.append(("grounded_sermon_editor", system_d, user_d))

    system_e = """You are a careful live transcript editor working from noisy translated source text.

Goal:
Produce the most faithful readable English update.

Decision rules:
1. NEW SOURCE TEXT is the primary evidence and should dominate the output.
2. Use PREVIOUS SOURCE TEXT and PREVIOUS REFINED CONTEXT only when they clearly help resolve or continue the new material.
3. If the new material starts a fresh thought, output only the new material.
4. Remove exact or near-exact repetition unless the repetition is clearly intentional rhetoric.
5. Do not add explanations, disclaimers, or meta commentary.
6. Prefer conservative wording over guessed meaning.
7. Return only the revised transcript text."""
    user_e = f"""Edit this transcript update.

[PREVIOUS REFINED CONTEXT]
"{prev_refined}"

[PREVIOUS SOURCE TEXT]
"{prev_raw}"

[NEW SOURCE TEXT]
"{new_raw}"
"""
    variants.append(("grounded_new_dominant", system_e, user_e))

    system_g = """You are a careful live transcript editor working from noisy translated source text.

Goal:
Produce the most faithful readable English update.

Decision rules:
1. NEW SOURCE TEXT is the primary evidence and should dominate the output.
2. Use PREVIOUS SOURCE TEXT and PREVIOUS REFINED CONTEXT only when they clearly help resolve or continue the new material.
3. The final answer must be English only.
4. If source text contains Spanish, Portuguese, French, or other non-English fragments, translate their meaning into English when reasonably clear.
5. Never copy non-English words into the output unless they are proper names.
6. If the new material starts a fresh thought, output only the new material.
7. Remove exact or near-exact repetition unless the repetition is clearly intentional rhetoric.
8. Do not add explanations, disclaimers, or meta commentary.
9. Prefer conservative wording over guessed meaning.
10. Return only the revised transcript text."""
    user_g = f"""Edit this transcript update.

[PREVIOUS REFINED CONTEXT]
"{prev_refined}"

[PREVIOUS SOURCE TEXT]
"{prev_raw}"

[NEW SOURCE TEXT]
"{new_raw}"
"""
    variants.append(("grounded_english_only_translate", system_g, user_g))

    system_h = """You are a careful live transcript editor working from noisy translated source text.

Goal:
Produce the most faithful readable English update.

Decision rules:
1. NEW SOURCE TEXT is the primary evidence and should dominate the output.
2. Use PREVIOUS SOURCE TEXT and PREVIOUS REFINED CONTEXT only when they clearly help resolve or continue the new material.
3. The final answer must be English only.
4. Treat any Spanish, Portuguese, French, or other non-English fragments as contamination from surrounding audio or translation noise unless they clearly repeat the same idea as nearby English.
5. When a non-English fragment clearly repeats nearby English content, keep only the English meaning once.
6. Never copy non-English words into the output unless they are proper names.
7. If the new material starts a fresh thought, output only the new material.
8. Remove exact or near-exact repetition unless the repetition is clearly intentional rhetoric.
9. Do not add explanations, disclaimers, or meta commentary.
10. Prefer conservative wording over guessed meaning.
11. Return only the revised transcript text."""
    user_h = f"""Edit this transcript update.

[PREVIOUS REFINED CONTEXT]
"{prev_refined}"

[PREVIOUS SOURCE TEXT]
"{prev_raw}"

[NEW SOURCE TEXT]
"{new_raw}"
"""
    variants.append(("grounded_english_only_ignore_noise", system_h, user_h))

    system_i = """You are a careful live transcript editor working from noisy translated source text.

Goal:
Produce the most faithful readable English update.

Decision rules:
1. NEW SOURCE TEXT is the primary evidence and should dominate the output.
2. Use PREVIOUS SOURCE TEXT and PREVIOUS REFINED CONTEXT only when they clearly help resolve or continue the new material.
3. The final answer must be English only.
4. Never copy non-English words into the output unless they are proper names.
5. If a non-English fragment appears as an isolated clause, trailing fragment, or side comment without a clear English anchor, drop it.
6. Only translate a non-English fragment when it is clearly central to the same idea and its meaning is reasonably obvious from context.
7. If the new material starts a fresh thought, output only the new material.
8. Remove exact or near-exact repetition unless the repetition is clearly intentional rhetoric.
9. Do not add explanations, disclaimers, or meta commentary.
10. Prefer conservative wording over guessed meaning.
11. Return only the revised transcript text."""
    user_i = f"""Edit this transcript update.

[PREVIOUS REFINED CONTEXT]
"{prev_refined}"

[PREVIOUS SOURCE TEXT]
"{prev_raw}"

[NEW SOURCE TEXT]
"{new_raw}"
"""
    variants.append(("grounded_english_only_drop_isolated", system_i, user_i))

    system_f = """You are a careful live transcript editor working from noisy translated source text.

Goal:
Produce the most faithful readable English update.

Decision rules:
1. NEW SOURCE TEXT is the primary evidence and should dominate the output.
2. Use PREVIOUS SOURCE TEXT and PREVIOUS REFINED CONTEXT only when they clearly help resolve or continue the new material.
3. If the new material starts a fresh thought, output only the new material.
4. Remove exact or near-exact repetition unless the repetition is clearly intentional rhetoric.
5. If the source repeats the same short clause three or more times in a row, keep at most two repetitions unless losing more would change the speaker's emphasis.
6. Do not add explanations, disclaimers, or meta commentary.
7. Prefer conservative wording over guessed meaning.
8. Return only the revised transcript text."""
    user_f = f"""Edit this transcript update.

[PREVIOUS REFINED CONTEXT]
"{prev_refined}"

[PREVIOUS SOURCE TEXT]
"{prev_raw}"

[NEW SOURCE TEXT]
"{new_raw}"
"""
    variants.append(("grounded_new_dominant_dedupe", system_f, user_f))

    if variant_filter:
        allowed = set(x.strip() for x in variant_filter.split(",") if x.strip())
        variants = [variant for variant in variants if variant[0] in allowed]
    return variants


def ollama_generate(model, system_prompt, user_prompt, url, timeout, keep_alive, temperature, top_p, repeat_penalty):
    payload = json.dumps(
        {
            "model": model,
            "prompt": user_prompt,
            "system": system_prompt,
            "stream": False,
            "options": {
                "temperature": temperature,
                "top_p": top_p,
                "repeat_penalty": repeat_penalty,
            },
            "keep_alive": keep_alive,
        }
    ).encode("utf-8")
    req = request.Request(
        url,
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with request.urlopen(req, timeout=timeout) as response:
            body = json.loads(response.read().decode("utf-8"))
    except error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"Ollama HTTP error {exc.code}: {detail}") from exc
    return body.get("response", "").strip()


def tokenize(text):
    return re.findall(r"[a-z0-9']+", (text or "").lower())


def repeated_ngram_ratio(text, n=3):
    tokens = tokenize(text)
    if len(tokens) < n:
        return 0.0
    grams = [tuple(tokens[i:i + n]) for i in range(len(tokens) - n + 1)]
    counts = Counter(grams)
    repeated = sum(count - 1 for count in counts.values() if count > 1)
    return repeated / max(1, len(grams))


def unsupported_token_ratio(output, evidence):
    evidence_tokens = set(tokenize(evidence))
    output_tokens = tokenize(output)
    if not output_tokens:
        return 0.0
    unsupported = [tok for tok in output_tokens if tok not in evidence_tokens]
    return len(unsupported) / len(output_tokens)


def evaluate_output(output, sample):
    evidence = " ".join(
        part for part in [sample["previous_raw"], sample["new_raw"]] if part
    )
    return {
        "output_len": len(tokenize(output)),
        "repeated_ngram_ratio": round(repeated_ngram_ratio(output), 4),
        "unsupported_token_ratio": round(unsupported_token_ratio(output, evidence), 4),
        "foreign_leak_ratio": round(foreign_leak_ratio(output), 4),
    }


def foreign_leak_ratio(output):
    foreign_markers = {
        "pero", "espanol", "tambien", "dijo", "dios", "responde",
        "mas", "ele", "continuou", "fe",
        "et", "rester", "ferme",
    }
    tokens = tokenize(output)
    if not tokens:
        return 0.0
    leaked = [tok for tok in tokens if tok in foreign_markers]
    return len(leaked) / len(tokens)


def parse_float_grid(raw_value, default_values):
    if not raw_value:
        return default_values
    return [float(x.strip()) for x in raw_value.split(",") if x.strip()]


def build_param_sets(args):
    temperatures = parse_float_grid(args.temperature_grid, [0.1])
    top_ps = parse_float_grid(args.top_p_grid, [0.9])
    repeat_penalties = parse_float_grid(args.repeat_penalty_grid, [1.1])

    param_sets = []
    for temperature in temperatures:
        for top_p in top_ps:
            for repeat_penalty in repeat_penalties:
                label = f"temp={temperature}_top_p={top_p}_repeat_penalty={repeat_penalty}"
                param_sets.append(
                    {
                        "label": label,
                        "temperature": temperature,
                        "top_p": top_p,
                        "repeat_penalty": repeat_penalty,
                    }
                )
    return param_sets


def run_experiments(args):
    samples = reconstruct_samples(args.log_file)
    samples = maybe_augment_samples(samples, args.sample_mode)
    if args.limit:
        samples = samples[: args.limit]

    param_sets = build_param_sets(args)
    results = []
    for index, sample in enumerate(samples):
        print(f"Running sample {index + 1}/{len(samples)}")
        sample_results = []
        for variant_name, system_prompt, user_prompt in build_variants(sample, args.variant_filter):
            print(f"  variant={variant_name}")
            for params in param_sets:
                print(f"    params={params['label']}")
                try:
                    output = ollama_generate(
                        args.model,
                        system_prompt,
                        user_prompt,
                        args.ollama_url,
                        args.timeout,
                        args.keep_alive,
                        params["temperature"],
                        params["top_p"],
                        params["repeat_penalty"],
                    )
                    metrics = evaluate_output(output, sample)
                    error_message = None
                except Exception as exc:
                    output = ""
                    metrics = None
                    error_message = str(exc)
                sample_results.append(
                    {
                        "variant": variant_name,
                        "params": params,
                        "output": output,
                        "metrics": metrics,
                        "error": error_message,
                    }
                )
        results.append(
            {
                "sample_index": index,
                "sample": sample,
                "results": sample_results,
            }
        )
    return results


def summarize_results(results):
    summary = {}
    for sample_result in results:
        for variant_result in sample_result["results"]:
            summary_key = f"{variant_result['variant']}|{variant_result['params']['label']}"
            bucket = summary.setdefault(
                summary_key,
                {
                    "repeated_ngram_ratio": [],
                    "unsupported_token_ratio": [],
                    "foreign_leak_ratio": [],
                    "output_len": [],
                    "errors": 0,
                },
            )
            if variant_result.get("error"):
                bucket["errors"] += 1
                continue
            for key, value in variant_result["metrics"].items():
                bucket[key].append(value)

    for variant_name, metrics in summary.items():
        summary[variant_name] = {
            key: round(statistics.mean(values), 4) if isinstance(values, list) and values else values
            for key, values in metrics.items()
        }
    return summary


def main():
    parser = argparse.ArgumentParser(description="Run prompt experiments against Ollama using distribute_debug.log samples.")
    parser.add_argument("--log-file", type=str, default="distribute_debug.log")
    parser.add_argument("--model", type=str, default="qwen2.5:3b-instruct")
    parser.add_argument("--ollama-url", type=str, default="http://127.0.0.1:11434/api/generate")
    parser.add_argument("--timeout", type=float, default=40.0)
    parser.add_argument("--keep-alive", type=str, default="30m")
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--output-file", type=str, default="logs/prompt_experiment_results.json")
    parser.add_argument("--variant-filter", type=str, default="", help="Comma-separated subset of variants to run.")
    parser.add_argument("--temperature-grid", type=str, default="", help="Comma-separated temperatures to test.")
    parser.add_argument("--top-p-grid", type=str, default="", help="Comma-separated top_p values to test.")
    parser.add_argument("--repeat-penalty-grid", type=str, default="", help="Comma-separated repeat_penalty values to test.")
    parser.add_argument("--sample-mode", type=str, default="base", help="Sample mode: base or multilingual_noise.")
    args = parser.parse_args()

    results = run_experiments(args)
    summary = summarize_results(results)
    payload = {"model": args.model, "summary": summary, "results": results}

    output_dir = os.path.dirname(args.output_file)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)
    with open(args.output_file, "w", encoding="utf-8") as handle:
        json.dump(payload, handle, ensure_ascii=False, indent=2)

    print(json.dumps(summary, ensure_ascii=False, indent=2))
    print(f"Wrote detailed results to {args.output_file}")


if __name__ == "__main__":
    main()