from __future__ import annotations

import csv
import json
import re
import shutil
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

DATA_INSTANCE_PREFIX = "data_"
DATA_SENTINEL = ".morehub_data.json"
SOURCE_MANIFEST = ".source_manifest.json"
_SELFIES_TOKEN_RE = re.compile(r"^(\[[^\[\]\s]+\])+$")


def sanitize_token(value: str) -> str:
    token = re.sub(r"[^A-Za-z0-9_.-]+", "_", value.strip())
    token = token.strip("._-")
    return token or "data"


def _sample_molecule_texts(path: Path, limit: int = 25) -> list[str]:
    samples: list[str] = []
    suffix = path.suffix.lower()

    try:
        if suffix == ".csv":
            with path.open("r", encoding="utf-8-sig", newline="") as handle:
                reader = csv.reader(handle)
                next(reader, None)  # header
                for row in reader:
                    if not row:
                        continue
                    text = str(row[0]).strip()
                    if not text:
                        continue
                    samples.append(text)
                    if len(samples) >= limit:
                        break
        else:
            with path.open("r", encoding="utf-8") as handle:
                for line in handle:
                    text = line.strip()
                    if not text:
                        continue
                    samples.append(text)
                    if len(samples) >= limit:
                        break
    except Exception:
        return samples

    return samples


def _looks_like_selfies(text: str) -> bool:
    value = text.strip()
    if not value:
        return False
    return bool(_SELFIES_TOKEN_RE.fullmatch(value))


def infer_molecule_kind(path: Path) -> str:
    name = path.name.lower()
    if "selfies" in name:
        return "selfies"
    if "smiles" in name:
        return "smiles"

    samples = _sample_molecule_texts(path)
    if samples:
        selfies_votes = sum(1 for sample in samples if _looks_like_selfies(sample))
        if selfies_votes >= max(1, int(0.6 * len(samples))):
            return "selfies"
        return "smiles"

    return "selfies"


def data_manifest_path(data_root: Path) -> Path:
    return data_root / DATA_SENTINEL


def source_manifest_path(data_root: Path) -> Path:
    return data_root / "data" / "source" / SOURCE_MANIFEST


def is_data_instance(data_root: Path) -> bool:
    return data_manifest_path(data_root).exists()


def _copy_or_move(src: Path, dst: Path, copy_only: bool) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    if copy_only:
        shutil.copy2(src, dst)
    else:
        shutil.move(str(src), str(dst))


def create_data_instance(
    workspace_root: Path,
    data_name: str,
    molecules_path: Path,
    target_path: Optional[Path],
    copy_only: bool,
) -> Path:
    safe_name = sanitize_token(data_name)
    data_root = workspace_root / f"{DATA_INSTANCE_PREFIX}{safe_name}"
    if data_root.exists():
        raise FileExistsError(f"Data folder already exists: {data_root}")

    source_dir = data_root / "data" / "source"
    embeddings_dir = data_root / "data" / "embeddings"
    benchmark_dir = data_root / "data" / "benchmarks" / "GP_regressor"

    source_dir.mkdir(parents=True, exist_ok=True)
    embeddings_dir.mkdir(parents=True, exist_ok=True)
    benchmark_dir.mkdir(parents=True, exist_ok=True)

    molecules_kind = infer_molecule_kind(molecules_path)
    molecules_ext = molecules_path.suffix.lower() or ".txt"
    source_type = "csv" if molecules_ext == ".csv" else "txt"

    molecules_filename = f"{safe_name}{molecules_ext}"
    molecules_dest = source_dir / molecules_filename
    _copy_or_move(molecules_path, molecules_dest, copy_only=copy_only)

    targets_filename = None
    targets_dest = None
    if target_path is not None:
        targets_ext = target_path.suffix.lower() or ".txt"
        targets_filename = f"targets{targets_ext}"
        targets_dest = source_dir / targets_filename
        _copy_or_move(target_path, targets_dest, copy_only=copy_only)

    created_at = datetime.now(timezone.utc).isoformat()

    source_manifest = {
        "format": molecules_kind,
        "source_type": source_type,
        "molecule_column": 0,
        "has_header": source_type == "csv",
        "molecules_file": molecules_filename,
        "targets_file": targets_filename,
        "molecules_original_name": molecules_path.name,
        "targets_original_name": target_path.name if target_path is not None else None,
        "molecules_original_path": str(molecules_path.resolve()),
        "targets_original_path": str(target_path.resolve()) if target_path is not None else None,
        "created_at": created_at,
    }
    source_manifest_path(data_root).write_text(json.dumps(source_manifest, indent=2) + "\n", encoding="utf-8")

    root_manifest = {
        "format": "morehub-data-instance",
        "version": 1,
        "name": safe_name,
        "created_at": created_at,
        "copy_only": copy_only,
        "source_manifest": str(source_manifest_path(data_root).relative_to(data_root)),
    }
    data_manifest_path(data_root).write_text(json.dumps(root_manifest, indent=2) + "\n", encoding="utf-8")

    return data_root


def load_data_manifest(data_root: Path) -> dict:
    path = data_manifest_path(data_root)
    if not path.exists():
        raise FileNotFoundError(f"Data manifest not found: {path}")
    return json.loads(path.read_text(encoding="utf-8"))


def load_source_manifest(data_root: Path) -> dict:
    path = source_manifest_path(data_root)
    if not path.exists():
        raise FileNotFoundError(f"Source manifest not found: {path}")
    return json.loads(path.read_text(encoding="utf-8"))


def resolve_data_source(data_root: Path) -> tuple[Path, Optional[Path], str]:
    source_dir = data_root / "data" / "source"
    manifest = load_source_manifest(data_root)

    molecules_file = manifest.get("molecules_file")
    if not molecules_file:
        raise FileNotFoundError(f"molecules_file missing in {source_manifest_path(data_root)}")

    molecules_path = source_dir / molecules_file
    if not molecules_path.exists():
        raise FileNotFoundError(f"Molecules file not found: {molecules_path}")

    targets_file = manifest.get("targets_file")
    target_path = (source_dir / targets_file) if targets_file else None
    if target_path is not None and not target_path.exists():
        raise FileNotFoundError(f"Targets file not found: {target_path}")

    input_kind = str(manifest.get("format") or infer_molecule_kind(molecules_path)).lower()
    return molecules_path, target_path, input_kind


def model_embeddings_dir(data_root: Path, model_name: str) -> Path:
    return data_root / "data" / "embeddings" / model_name.upper()


def embeddings_manifest_path(model_embeddings_dir_path: Path) -> Path:
    return model_embeddings_dir_path / ".embeddings_manifest.json"


def write_embeddings_manifest(model_embeddings_dir_path: Path, payload: dict) -> None:
    model_embeddings_dir_path.mkdir(parents=True, exist_ok=True)
    embeddings_manifest_path(model_embeddings_dir_path).write_text(
        json.dumps(payload, indent=2) + "\n",
        encoding="utf-8",
    )


def find_data_root(start_path: Path) -> Optional[Path]:
    path = start_path.resolve()
    for candidate in [path, *path.parents]:
        if data_manifest_path(candidate).exists():
            return candidate
    return None

