import argparse
import os
import sys
from pathlib import Path
from typing import List, Tuple

import numpy as np
import torch
import torch.nn.functional as F

SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_HF_HOME = SCRIPT_DIR / ".cache" / "huggingface"
DEFAULT_TRANSFORMERS_CACHE = DEFAULT_HF_HOME / "transformers"
DEFAULT_LOCAL_MODEL_DIR = DEFAULT_HF_HOME / "MoLFormer-XL-both-10pct"

# Ensure huggingface cache is model-local by default.
os.environ.setdefault("HF_HOME", str(DEFAULT_HF_HOME))
os.environ.setdefault("TRANSFORMERS_CACHE", str(DEFAULT_TRANSFORMERS_CACHE))
os.environ.setdefault("HF_HUB_CACHE", str(DEFAULT_HF_HOME / "hub"))
os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")
os.environ.setdefault("HF_HUB_DISABLE_IMPLICIT_TOKEN", "1")

try:
    from transformers import AutoModel, AutoTokenizer  # type: ignore
except Exception as exc:
    raise ImportError(
        "Could not import transformers AutoModel/AutoTokenizer. "
        "Ensure MOLFORMER setup completed successfully."
    ) from exc


PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from embeddings_preprocess import (
    load_input_lines,
    preprocess_lines,
    resolve_output_paths,
    write_failures_csv,
    write_mapping_csv,
)


def _resolve_device(device_arg: str) -> torch.device:
    normalized = device_arg.strip().lower()
    if normalized == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if normalized not in {"cpu", "cuda"}:
        raise ValueError(f"Unsupported device: {device_arg}")
    if normalized == "cuda" and not torch.cuda.is_available():
        print("CUDA requested but unavailable; falling back to CPU.", file=sys.stderr)
        return torch.device("cpu")
    return torch.device(normalized)


def _resolve_model_source(model_path: str, pretrained_name: str) -> Tuple[str, bool]:
    explicit = model_path.strip()
    if explicit:
        candidate = Path(explicit).expanduser()
        if candidate.exists():
            return str(candidate.resolve()), True
        return explicit, False

    if DEFAULT_LOCAL_MODEL_DIR.exists():
        return str(DEFAULT_LOCAL_MODEL_DIR.resolve()), True

    return pretrained_name, False


def _build_model(
    model_source: str,
    revision: str,
    trust_remote_code: bool,
    local_files_only: bool,
    device: torch.device,
):
    tokenizer_kwargs = {
        "trust_remote_code": trust_remote_code,
        "local_files_only": local_files_only,
    }
    if revision:
        tokenizer_kwargs["revision"] = revision

    tokenizer = AutoTokenizer.from_pretrained(
        model_source,
        **tokenizer_kwargs,
    )

    model_kwargs = {
        "trust_remote_code": trust_remote_code,
        "local_files_only": local_files_only,
    }
    if revision:
        model_kwargs["revision"] = revision

    try:
        model = AutoModel.from_pretrained(
            model_source,
            deterministic_eval=True,
            **model_kwargs,
        )
    except TypeError:
        model = AutoModel.from_pretrained(
            model_source,
            **model_kwargs,
        )

    return tokenizer, model.to(device).eval()


@torch.inference_mode()
def _encode_batches(
    model_inputs: List[str],
    model,
    tokenizer,
    device: torch.device,
    batch_size: int,
    max_length: int,
    normalize: bool,
) -> np.ndarray:
    chunks = []

    for start in range(0, len(model_inputs), batch_size):
        batch_texts = model_inputs[start : start + batch_size]
        encoded = tokenizer(
            batch_texts,
            padding=True,
            truncation=True,
            max_length=max_length,
            return_tensors="pt",
        )
        encoded = {key: value.to(device) for key, value in encoded.items()}
        outputs = model(**encoded)

        pooler = getattr(outputs, "pooler_output", None)
        if pooler is not None:
            emb = pooler
        else:
            hidden = getattr(outputs, "last_hidden_state", None)
            if hidden is None:
                raise RuntimeError("MOLFORMER outputs did not include pooler_output or last_hidden_state.")
            mask = encoded.get("attention_mask")
            if mask is None:
                mask = torch.ones(hidden.shape[:2], dtype=torch.long, device=device)
            mask = mask.unsqueeze(-1).type_as(hidden)
            emb = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-6)

        if normalize:
            emb = F.normalize(emb, dim=-1)

        chunks.append(emb.detach().cpu().numpy().astype(np.float32, copy=False))

    if not chunks:
        raise RuntimeError("No embeddings were produced by MOLFORMER.")
    return np.concatenate(chunks, axis=0)


def main() -> None:
    parser = argparse.ArgumentParser(description="Encode molecules with MOLFORMER and save embeddings")
    parser.add_argument("--file", required=True, help="Path to input .txt file")
    parser.add_argument(
        "--input-format",
        default="selfies",
        choices=["selfies", "smiles"],
        help="Format of rows in input file",
    )
    parser.add_argument(
        "--model-input-type",
        default="smiles",
        choices=["selfies", "smiles"],
        help="Expected model input type from metadata",
    )
    parser.add_argument("--output-dir", default=".", help="Directory to write outputs")
    parser.add_argument("--output-stem", default="embeddings", help="Base output filename stem")
    parser.add_argument(
        "--pretrained-name",
        default="ibm/MoLFormer-XL-both-10pct",
        help="Hugging Face repo id or local model path",
    )
    parser.add_argument(
        "--model-path",
        default="",
        help="Optional explicit local model directory (overrides default cache lookup)",
    )
    parser.add_argument(
        "--revision",
        default="",
        help="Optional model revision override",
    )
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--max-length", type=int, default=256)
    parser.add_argument(
        "--device",
        default="auto",
        choices=["auto", "cpu", "cuda"],
        help="Inference device",
    )
    parser.add_argument(
        "--trust-remote-code",
        dest="trust_remote_code",
        action="store_true",
        help="Allow custom model code from the Hugging Face repository",
    )
    parser.add_argument(
        "--no-trust-remote-code",
        dest="trust_remote_code",
        action="store_false",
        help="Disable remote model code loading",
    )
    parser.add_argument(
        "--no-normalize",
        dest="normalize",
        action="store_false",
        help="Disable L2-normalization before saving",
    )
    parser.set_defaults(normalize=True, trust_remote_code=True)
    args = parser.parse_args()

    input_path = Path(args.file).resolve()
    if not input_path.exists():
        print(f"Input file not found: {input_path}", file=sys.stderr)
        sys.exit(1)

    output_dir = Path(args.output_dir).resolve()
    npy_path, map_path, fail_path = resolve_output_paths(output_dir, args.output_stem)

    lines = load_input_lines(input_path)
    records, failures = preprocess_lines(lines, args.input_format, args.model_input_type)
    if not records:
        print("No valid molecules after circularity validation. Aborting.", file=sys.stderr)
        if failures:
            print(f"Failures: {len(failures)} (see {fail_path})", file=sys.stderr)
            write_failures_csv(fail_path, failures)
        sys.exit(1)

    model_inputs = [record["model_input"] for record in records]

    model_source, local_files_only = _resolve_model_source(args.model_path, args.pretrained_name)
    device = _resolve_device(args.device)
    tokenizer, model = _build_model(
        model_source=model_source,
        revision=args.revision.strip(),
        trust_remote_code=args.trust_remote_code,
        local_files_only=local_files_only,
        device=device,
    )

    embeddings = _encode_batches(
        model_inputs=model_inputs,
        model=model,
        tokenizer=tokenizer,
        device=device,
        batch_size=args.batch_size,
        max_length=args.max_length,
        normalize=args.normalize,
    )

    np.save(npy_path, embeddings)
    write_mapping_csv(map_path, records)
    if failures:
        write_failures_csv(fail_path, failures)

    print(f"Saved embeddings to {npy_path} with shape {embeddings.shape}.")
    print(f"Saved row mapping to {map_path}.")
    print(f"Model source: {model_source} (local_files_only={local_files_only})")
    if args.revision.strip():
        print(f"Revision: {args.revision.strip()}")
    if failures:
        print(f"Filtered {len(failures)} invalid rows. See {fail_path}.")


if __name__ == "__main__":
    main()