import nltk
import datetime
from pathlib import Path
import logging
import sys
import json
import shutil
from typing import List

logger = logging.getLogger(__name__)

def set_from_json(path: Path) -> set[str]:
    """Return a hashable set of strings from a ban-list file that may be
       either  ["foo bar", …]  or  [["foo bar", 1], …]."""
    if not path or not path.is_file():
        return set()
    try:
        raw = json.loads(path.read_text("utf-8"))
    except Exception:
        return set()

    if not isinstance(raw, list):
        return set()

    out = set()
    for item in raw:
        if isinstance(item, list) and item:        # slop format
            out.add(str(item[0]))
        else:                                      # plain string
            out.add(str(item))
    return out


def merge_custom_bans_into_file(path: Path, extra_items: List[str]) -> None:
    """Merge extra_items into `path`, preserving original on-disk format."""

    # 1) read whatever is already there
    try:
        current_raw = json.loads(path.read_text("utf-8")) if path.exists() else []
    except Exception:
        current_raw = []

    if not isinstance(current_raw, list):
        current_raw = []

    # 2) normalise existing items → plain strings
    existing: set[str] = set()
    slop_format = False               # do we need to write back [[phrase,1]] ?

    for entry in current_raw:
        if isinstance(entry, list):   # slop-phrase style [phrase, freq]
            slop_format = True
            if entry:                 # non-empty list
                existing.add(str(entry[0]))
        else:                         # plain string
            existing.add(str(entry))

    # 3) merge & sort
    merged = sorted(existing | set(map(str, extra_items)))

    # 4) write back in the same shape we found
    if slop_format:
        payload = [[p, 1] for p in merged]
    else:
        payload = merged

    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, ensure_ascii=False), "utf-8")


def download_nltk_resource(resource_id: str, resource_name: str):
    """Downloads NLTK resource if not found."""
    try:
        nltk.data.find(resource_id)
        logger.debug(f"NLTK '{resource_name}' resource found.")
    except LookupError:
        logger.info(f"NLTK '{resource_name}' resource not found. Downloading...")
        try:
            nltk.download(resource_name, quiet=True)
            logger.info(f"NLTK '{resource_name}' resource downloaded.")
        except Exception as e:
            logger.warning(f"Could not automatically download NLTK '{resource_name}' resource: {e}. "
                           "Manual download might be required (e.g., python -m nltk.downloader punkt stopwords)")
    except Exception as e:
        logger.warning(f"Error checking NLTK '{resource_name}' resource: {e}.")

def create_experiment_dir(base_dir_path: Path, resume_dir: Path | None = None) -> Path:
    """
    Determine the directory that the pipeline should work in.

    • When --resume-from-dir is supplied we MUST use that exact path.
      If the directory is missing or not a directory, raise immediately.

    • When no resume dir is given, create a new timestamped directory under
      *base_dir_path* (parents created as needed) and return it.
    """
    if resume_dir is not None:
        if resume_dir.is_dir():
            logger.info(f"Resuming experiment in existing directory: {resume_dir.resolve()}")
            return resume_dir
        # hard-fail: the user explicitly asked to resume here
        raise FileNotFoundError(
            f"--resume-from-dir was set to '{resume_dir}', "
            "but that path does not exist or is not a directory."
        )

    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    experiment_dir = base_dir_path / f"run_{timestamp}"
    experiment_dir.mkdir(parents=True, exist_ok=False)
    logger.info(f"Created experiment directory: {experiment_dir.resolve()}")
    return experiment_dir


def ensure_antislop_vllm_config_exists(antislop_vllm_dir: Path):
    """
    Copies antislop-vllm/config-example.yaml to config.yaml if config.yaml is absent.
    This is a helper for users, but the main pipeline will pass params via CLI.
    """
    cfg_path = antislop_vllm_dir / "config.yaml"
    example_path = antislop_vllm_dir / "config-example.yaml"

    if not cfg_path.exists():
        if example_path.exists():
            try:
                shutil.copy(example_path, cfg_path)
                logger.info(f"Copied {example_path} to {cfg_path} for antislop-vllm (user convenience).")
            except Exception as e:
                logger.warning(f"Could not copy antislop-vllm config example: {e}")
        else:
            logger.debug("antislop-vllm/config-example.yaml not found. No default config.yaml created for it.")
    else:
        logger.debug("antislop-vllm/config.yaml already exists.")


###############################################################################
# NLTK helpers
###############################################################################
CORE_NLTK_RESOURCES = [
    ("tokenizers/punkt",       "punkt"),        # sentence + word tokeniser data
    ("tokenizers/punkt_tab",   "punkt_tab"),    # new in NLTK 3.9+, used by PunktTokenizer
    ("corpora/stopwords",      "stopwords"),    # obvious
]

def ensure_core_nltk_resources() -> None:
    """
    Download the three NLTK resources our pipeline needs *once* at start-up.
    Safe to call multiple times – it‘s a no-op if they’re already present.
    """
    for resource_id, resource_name in CORE_NLTK_RESOURCES:
        download_nltk_resource(resource_id, resource_name)