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

import numpy as np
import selfies as sf
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"

MODEL_VARIANTS: Dict[str, Dict[str, str]] = {
    "7b": {
        "repo_id": "zjunlp/MolGen-7b",
        "cache_subdir": "MolGen-7b",
    },
    "large": {
        "repo_id": "zjunlp/MolGen-large",
        "cache_subdir": "MolGen-large",
    },
    "large-opt": {
        "repo_id": "zjunlp/MolGen-large-opt",
        "cache_subdir": "MolGen-large-opt",
    },
}

# 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 AutoConfig, AutoModel, AutoTokenizer  # type: ignore
except Exception as exc:
    raise ImportError(
        "Could not import transformers AutoConfig/AutoModel/AutoTokenizer. "
        "Ensure MOLGEN 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 _normalize_variant(variant: str) -> str:
    normalized = variant.strip().lower()
    if normalized not in MODEL_VARIANTS:
        allowed = ", ".join(sorted(MODEL_VARIANTS.keys()))
        raise ValueError(f"Unsupported MOLGEN variant: {variant}. Allowed: {allowed}")
    return normalized


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, variant: 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

    explicit_pretrained = pretrained_name.strip()
    if explicit_pretrained:
        candidate = Path(explicit_pretrained).expanduser()
        if candidate.exists():
            return str(candidate.resolve()), True
        return explicit_pretrained, False

    variant_cfg = MODEL_VARIANTS[variant]
    local_dir = DEFAULT_HF_HOME / variant_cfg["cache_subdir"]
    if local_dir.exists():
        return str(local_dir.resolve()), True

    return variant_cfg["repo_id"], False


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

    config = AutoConfig.from_pretrained(model_source, **kwargs)
    tokenizer = AutoTokenizer.from_pretrained(model_source, **kwargs)
    model = AutoModel.from_pretrained(model_source, **kwargs).to(device).eval()

    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
        tokenizer.pad_token = tokenizer.eos_token

    return config, tokenizer, model


def _to_molgen_text(canonical_smiles: str) -> str:
    selfies_text = sf.encoder(canonical_smiles)
    return selfies_text.replace("][", "] [")


@torch.inference_mode()
def _encode_batches(
    model_inputs: List[str],
    config,
    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()}

        if bool(getattr(config, "is_encoder_decoder", False)) and hasattr(model, "get_encoder"):
            encoder_outputs = model.get_encoder()(
                input_ids=encoded["input_ids"],
                attention_mask=encoded.get("attention_mask"),
            )
            last_hidden = encoder_outputs.last_hidden_state
        else:
            outputs = model(
                input_ids=encoded["input_ids"],
                attention_mask=encoded.get("attention_mask"),
            )
            last_hidden = getattr(outputs, "last_hidden_state", None)
            if last_hidden is None:
                if isinstance(outputs, (tuple, list)) and outputs:
                    last_hidden = outputs[0]
                else:
                    raise RuntimeError("MOLGEN forward outputs did not include last_hidden_state.")

        mask = encoded.get("attention_mask")
        if mask is None:
            mask = torch.ones(last_hidden.shape[:2], dtype=torch.long, device=device)
        mask = mask.unsqueeze(-1).type_as(last_hidden)
        emb = (last_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 MOLGEN.")
    return np.concatenate(chunks, axis=0)


def main() -> None:
    parser = argparse.ArgumentParser(description="Encode molecules with MOLGEN 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(
        "--variant",
        default="7b",
        choices=sorted(MODEL_VARIANTS.keys()),
        help="MolGen checkpoint variant",
    )
    parser.add_argument(
        "--pretrained-name",
        default="",
        help="Override model source (HF repo id or local model path)",
    )
    parser.add_argument(
        "--model-path",
        default="",
        help="Optional explicit local model directory (highest priority)",
    )
    parser.add_argument(
        "--revision",
        default="",
        help="Optional model revision override",
    )
    parser.add_argument("--batch-size", type=int, default=8)
    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)

    molgen_inputs: List[str] = []
    filtered_records = []
    for record in records:
        try:
            molgen_inputs.append(_to_molgen_text(record["smiles"]))
            filtered_records.append(record)
        except Exception:
            failures.append(
                {
                    "original_line": record["original_line"],
                    "input_text": record["input_text"],
                    "reason": "selfies_encoding_failed",
                }
            )

    if not molgen_inputs:
        print("No valid molecules after MolGen token preparation. Aborting.", file=sys.stderr)
        if failures:
            write_failures_csv(fail_path, failures)
        sys.exit(1)

    variant = _normalize_variant(args.variant)
    model_source, local_files_only = _resolve_model_source(args.model_path, args.pretrained_name, variant)
    device = _resolve_device(args.device)

    config, 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=molgen_inputs,
        config=config,
        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, filtered_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"MolGen variant: {variant}")
    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()