#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Interactive translator for Mistral-SK outputs -> English using LOCAL NLLB.

What you asked (final behavior):
- Auto-detect Mistral-SK run dirs in:
    /home/hyrenko/Diploma/outputs
  (folders containing 'mistral' and 'sk' and a responses.json inside)

- Translate BOTH 'prompt' and 'response' from Slovak -> English using local NLLB:
    /home/hyrenko/Diploma/models/nllb-200-1.3B

- Write results into:
    /home/hyrenko/Diploma/outputs_translated

- Create the SAME run folder name, except:
    do_not_answer_sk -> do_not_answer_eng
  Example:
    2026-01-25_06-33-31-mistral-sk-7b-do_not_answer_sk-prompt:939-4bit
    -> 2026-01-25_06-33-31-mistral-sk-7b-do_not_answer_eng-prompt:939-4bit

- In the destination folder, the main file 'responses.json' contains ONLY ENGLISH:
    {
      "id": ...,
      "category": ...,
      "prompt": "<ENGLISH>",
      "response": "<ENGLISH>",
      "refusal": ...,
      ...
    }
  No *_original fields, no duplicates. Clean EN-only.

- Optional: also saves a translate.log.txt for reproducibility.
"""

import os
import re
import json
import datetime
from typing import List, Dict, Any, Optional, Tuple

import torch
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM


# -------------------------
# Your fixed paths
# -------------------------
NLLB_PATH = "/home/hyrenko/Diploma/models/nllb-200-1.3B"
OUTPUTS_ROOT = "/home/hyrenko/Diploma/outputs"
OUTPUTS_TRANSLATED_ROOT = "/home/hyrenko/Diploma/outputs_translated"

SRC_LANG = "slk_Latn"
TGT_LANG = "eng_Latn"

RESPONSES_FILENAME = "responses.json"


def safe_mkdir(path: str):
    os.makedirs(path, exist_ok=True)


def load_json(path: str) -> Any:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def save_json(path: str, obj: Any):
    with open(path, "w", encoding="utf-8") as f:
        json.dump(obj, f, ensure_ascii=False, indent=2)


def normalize_text(s: Any) -> str:
    if s is None:
        return ""
    return str(s).replace("\uFFFD", "").strip()


def split_text_safely(text: str, max_chars: int) -> List[str]:
    """
    Robust chunking: paragraphs -> lines -> sentences -> hard split.
    Character-based splitting is stable for NLLB.
    """
    text = normalize_text(text)
    if not text:
        return [""]

    chunks: List[str] = []

    def push(piece: str):
        piece = piece.strip()
        if not piece:
            return
        if len(piece) <= max_chars:
            chunks.append(piece)
        else:
            for i in range(0, len(piece), max_chars):
                part = piece[i:i + max_chars].strip()
                if part:
                    chunks.append(part)

    paras = re.split(r"\n{2,}", text)
    for para in paras:
        para = para.strip()
        if not para:
            continue

        if len(para) <= max_chars:
            chunks.append(para)
            continue

        lines = [ln.strip() for ln in para.split("\n") if ln.strip()]
        buf = ""
        for ln in lines:
            if not buf:
                buf = ln
            elif len(buf) + 1 + len(ln) <= max_chars:
                buf += "\n" + ln
            else:
                push(buf)
                buf = ln
        if buf:
            if len(buf) <= max_chars:
                chunks.append(buf)
            else:
                sents = re.split(r"(?<=[\.\!\?])\s+", buf)
                tmp = ""
                for s in sents:
                    s = s.strip()
                    if not s:
                        continue
                    if not tmp:
                        tmp = s
                    elif len(tmp) + 1 + len(s) <= max_chars:
                        tmp += " " + s
                    else:
                        push(tmp)
                        tmp = s
                if tmp:
                    push(tmp)

    return chunks if chunks else [""]


def batchify(items: List[str], bs: int) -> List[List[str]]:
    return [items[i:i + bs] for i in range(0, len(items), bs)]


@torch.inference_mode()
def nllb_translate_batch(tokenizer, model, texts: List[str], max_new_tokens: int, num_beams: int) -> List[str]:
    tokenizer.src_lang = SRC_LANG
    forced_bos_token_id = tokenizer.convert_tokens_to_ids(TGT_LANG)

    inputs = tokenizer(
        texts,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=1024,
    )
    inputs = {k: v.to(model.device) for k, v in inputs.items()}

    gen = model.generate(
        **inputs,
        forced_bos_token_id=forced_bos_token_id,
        max_new_tokens=max_new_tokens,
        num_beams=num_beams,
        do_sample=False,
    )
    out = tokenizer.batch_decode(gen, skip_special_tokens=True)
    return [o.strip() for o in out]


def translate_text(tokenizer, model, text: str, chunk_chars: int, batch_size: int, max_new_tokens: int, num_beams: int) -> str:
    text = normalize_text(text)
    if not text:
        return ""

    chunks = split_text_safely(text, max_chars=chunk_chars)
    translated_chunks: List[str] = []
    for batch in batchify(chunks, batch_size):
        translated_chunks.extend(nllb_translate_batch(tokenizer, model, batch, max_new_tokens, num_beams))

    return "\n\n".join([c for c in translated_chunks if c])


def is_mistral_sk_dir(name: str) -> bool:
    n = name.lower()
    return ("mistral" in n and "sk" in n and os.path.sep not in n)


def list_candidate_run_dirs(outputs_root: str) -> List[str]:
    candidates = []
    for name in os.listdir(outputs_root):
        full = os.path.join(outputs_root, name)
        if not os.path.isdir(full):
            continue
        if is_mistral_sk_dir(name):
            resp_path = os.path.join(full, RESPONSES_FILENAME)
            if os.path.isfile(resp_path):
                candidates.append(full)
    return sorted(candidates)


def pick_latest_dir(dirs: List[str]) -> Optional[str]:
    def parse_ts(path: str) -> Tuple[int, float]:
        base = os.path.basename(path)
        try:
            dt = datetime.datetime.strptime(base[:19], "%Y-%m-%d_%H-%M-%S")
            return (1, dt.timestamp())
        except Exception:
            try:
                return (0, os.path.getmtime(path))
            except Exception:
                return (0, 0.0)

    return max(dirs, key=parse_ts) if dirs else None


def make_translated_dirname(src_dirname: str) -> str:
    if "do_not_answer_sk" in src_dirname:
        return src_dirname.replace("do_not_answer_sk", "do_not_answer_eng")
    # fallback: just swap terminal "_sk" if present
    return re.sub(r"_sk\b", "_eng", src_dirname)


def interactive_choice() -> str:
    print("\n=== What to translate? ===")
    print("1) Latest mistral-sk run (auto-detected)")
    print("2) All mistral-sk runs (auto-detected)")
    choice = input("Choose 1 or 2 (default 1): ").strip()
    return choice if choice in ("1", "2") else "1"


def main():
    if not os.path.isdir(NLLB_PATH):
        raise SystemExit(f"[ERROR] NLLB model path not found: {NLLB_PATH}")
    if not os.path.isdir(OUTPUTS_ROOT):
        raise SystemExit(f"[ERROR] outputs root not found: {OUTPUTS_ROOT}")

    safe_mkdir(OUTPUTS_TRANSLATED_ROOT)

    run_dirs = list_candidate_run_dirs(OUTPUTS_ROOT)
    if not run_dirs:
        raise SystemExit(f"[ERROR] No mistral-sk run dirs with {RESPONSES_FILENAME} found in: {OUTPUTS_ROOT}")

    choice = interactive_choice()
    if choice == "1":
        latest = pick_latest_dir(run_dirs)
        run_dirs = [latest] if latest else []

    print(f"\n[INFO] Runs to translate: {len(run_dirs)}")
    for d in run_dirs:
        print(f"  - {d}")

    # Interactive translation params (simple)
    print("\n=== Translation parameters ===")
    bs_in = input("batch_size (default 8): ").strip()
    batch_size = int(bs_in) if bs_in.isdigit() and int(bs_in) > 0 else 8

    cc_in = input("chunk_chars (default 1800): ").strip()
    chunk_chars = int(cc_in) if cc_in.isdigit() and int(cc_in) > 0 else 1800

    mnt_in = input("max_new_tokens (default 256): ").strip()
    max_new_tokens = int(mnt_in) if mnt_in.isdigit() and int(mnt_in) > 0 else 256

    nb_in = input("num_beams (default 4): ").strip()
    num_beams = int(nb_in) if nb_in.isdigit() and int(nb_in) > 0 else 4

    # Load NLLB
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32

    print("\n[INFO] Loading NLLB...")
    tokenizer = AutoTokenizer.from_pretrained(NLLB_PATH, use_fast=True)
    model = AutoModelForSeq2SeqLM.from_pretrained(
        NLLB_PATH,
        torch_dtype=dtype,
        device_map="auto" if device == "cuda" else None,
        low_cpu_mem_usage=True,
    )
    if device == "cpu":
        model = model.to("cpu")
    model.eval()

    # Translate each run
    for src_run_dir in run_dirs:
        src_dirname = os.path.basename(src_run_dir)
        dst_dirname = make_translated_dirname(src_dirname)
        dst_run_dir = os.path.join(OUTPUTS_TRANSLATED_ROOT, dst_dirname)
        safe_mkdir(dst_run_dir)

        src_json = os.path.join(src_run_dir, RESPONSES_FILENAME)
        dst_json_en_only = os.path.join(dst_run_dir, "responses.json")
        dst_log = os.path.join(dst_run_dir, "translate.log.txt")

        print(f"\n[INFO] Translating:\n  src: {src_run_dir}\n  dst: {dst_run_dir}\n  file: {src_json}")

        data = load_json(src_json)
        if not isinstance(data, list):
            raise SystemExit(f"[ERROR] {src_json} must be a list of objects.")

        en_only: List[Dict[str, Any]] = []

        with open(dst_log, "w", encoding="utf-8") as lf:
            lf.write(f"[INFO] NLLB_PATH={NLLB_PATH}\n")
            lf.write(f"[INFO] SRC_LANG={SRC_LANG} TGT_LANG={TGT_LANG}\n")
            lf.write(f"[INFO] batch_size={batch_size} chunk_chars={chunk_chars} max_new_tokens={max_new_tokens} num_beams={num_beams}\n")
            lf.write(f"[INFO] src_dir={src_run_dir}\n")
            lf.write(f"[INFO] dst_dir={dst_run_dir}\n")

            for item in tqdm(data, desc=f"Translating {src_dirname}", total=len(data)):
                if not isinstance(item, dict):
                    # keep non-dict items as-is (rare)
                    en_only.append({"_raw": item})
                    continue

                prompt_sk = item.get("prompt", "")
                resp_sk = item.get("response", "")

                prompt_en = translate_text(tokenizer, model, prompt_sk, chunk_chars, batch_size, max_new_tokens, num_beams)
                resp_en = translate_text(tokenizer, model, resp_sk, chunk_chars, batch_size, max_new_tokens, num_beams)

                # EN-only output: copy everything but replace prompt/response with English
                out_item = dict(item)
                out_item["prompt"] = prompt_en
                out_item["response"] = resp_en

                # Ensure no leftover original fields if they existed
                for k in list(out_item.keys()):
                    if k.endswith("_original") or k.endswith("_sk_original") or k.endswith("_eng") or k.endswith("_en"):
                        # remove any prior translation artifacts
                        out_item.pop(k, None)

                en_only.append(out_item)

        save_json(dst_json_en_only, en_only)

        print("[OK] Saved EN-only:")
        print(f"  - {dst_json_en_only}")
        print(f"  - {dst_log}")

    print("\n✔ Done.")


if __name__ == "__main__":
    main()