import argparse
import inspect
import sys
from pathlib import Path

import numpy as np
import torch

try:
    import xfmr_compat

    xfmr_compat.apply()
except Exception:
    pass

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 safe_load_ckpt(path):
    try:
        from pytorch_lightning.utilities.parsing import AttributeDict as PLAttrDict

        torch.serialization.add_safe_globals([PLAttrDict])
    except Exception:
        pass
    try:
        return torch.load(path, map_location="cpu", weights_only=False)
    except TypeError:
        return torch.load(path, map_location="cpu")


def hp_to_namespace(hp_dict):
    from argparse import Namespace

    return Namespace(**hp_dict)


def import_obj(path):
    if ":" in path:
        mod, obj = path.split(":", 1)
    else:
        parts = path.split(".")
        mod, obj = ".".join(parts[:-1]), parts[-1]
    import importlib

    module = importlib.import_module(mod)
    return getattr(module, obj)


class FeaturizerTokenizerAdapter:
    def __init__(self, featurizer, ns):
        self.featurizer = featurizer
        self.ns = ns
        self.token_to_idx = getattr(self.featurizer, "token_to_idx", None)
        if self.token_to_idx is None:
            raise RuntimeError("Featurizer missing token_to_idx; cannot map tokens to ids.")
        self.special_ids = self._build_special_ids()

    def _find_token_id(self, candidates):
        for candidate in candidates:
            if candidate in self.token_to_idx:
                return self.token_to_idx[candidate]
        return None

    def _build_special_ids(self):
        specials = {
            "pad": ["[PAD]", "<pad>", "<PAD>", "PAD"],
            "cls": ["[CLS]", "<cls>", "<CLS>", "CLS", "<bos>", "[BOS]"],
            "sep": ["[SEP]", "<sep>", "<SEP>", "SEP", "<eos>", "[EOS]"],
        }
        ids = {name: self._find_token_id(tokens) for name, tokens in specials.items()}
        if ids["pad"] is None:
            ids["pad"] = 0
        if ids["cls"] is None:
            ids["cls"] = 1
        if ids["sep"] is None:
            ids["sep"] = 2
        return ids

    def _encode_one(self, text):
        try:
            tokens = self.featurizer.encode(text)
        except TypeError:
            tokens = self.featurizer.encode(text=text)
        if isinstance(tokens, str):
            tokens = list(tokens)
        unknown_token = self.token_to_idx.get("[UNK]", self.special_ids["pad"])
        return [self.token_to_idx.get(token, unknown_token) for token in tokens]

    def __call__(self, batch_texts, padding=True, truncation=True, max_length=128, return_tensors="pt"):
        del return_tensors
        input_ids, attention_masks = [], []
        cls_id = self.special_ids["cls"]
        sep_id = self.special_ids["sep"]
        pad_id = self.special_ids["pad"]

        for text in batch_texts:
            ids = [cls_id] + self._encode_one(text) + [sep_id]
            if truncation and len(ids) > max_length:
                ids = ids[:max_length]
            mask = [1] * len(ids)
            if padding and len(ids) < max_length:
                pad_count = max_length - len(ids)
                ids = ids + [pad_id] * pad_count
                mask = mask + [0] * pad_count
            input_ids.append(ids)
            attention_masks.append(mask)

        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attention_masks, dtype=torch.long),
        }


def try_build_featurizer_adapter(ns):
    try:
        import importlib
        import inspect as inspect_module

        module = importlib.import_module("molbert.utils.featurizer.molfeaturizer")
    except Exception:
        return None, "Could not import molbert.utils.featurizer.molfeaturizer"

    for _, obj in inspect_module.getmembers(module, inspect_module.isclass):
        if obj.__module__ != module.__name__:
            continue
        try:
            instance = obj(ns)
        except Exception:
            try:
                instance = obj()
            except Exception:
                continue
        if not (hasattr(instance, "encode") or hasattr(instance, "__call__")):
            continue
        if getattr(instance, "token_to_idx", None) is None:
            continue
        return FeaturizerTokenizerAdapter(instance, ns), None

    return None, "No suitable featurizer with encode/token_to_idx found"


def patch_embeddings_forward(backbone):
    try:
        embeddings = backbone.embeddings
    except Exception:
        return

    original_forward = embeddings.forward

    def patched_forward(*args, **kwargs):
        kwargs.pop("past_key_values_length", None)
        return original_forward(*args, **kwargs)

    embeddings.forward = patched_forward


def find_backbone(model):
    if hasattr(model, "model") and hasattr(model.model, "bert"):
        return model.model.bert
    if hasattr(model, "bert"):
        return model.bert
    if hasattr(model, "model") and hasattr(model.model, "backbone"):
        return model.model.backbone
    if hasattr(model, "backbone"):
        return model.backbone
    return None


@torch.no_grad()
def batch_embed_inputs(model_inputs, model, device, batch_size, max_length, use_cls_embed):
    tokenizer = getattr(model, "tokenizer", None)
    if tokenizer is None:
        raise RuntimeError("Model has no `.tokenizer`.")

    model.eval().to(device)
    backbone = find_backbone(model)
    if backbone is None:
        raise RuntimeError("Could not locate a raw BERT-like backbone on the model.")
    patch_embeddings_forward(backbone)

    signature = inspect.signature(backbone.forward)
    wanted_args = set(signature.parameters.keys())

    def filter_kwargs(**kwargs):
        return {name: value for name, value in kwargs.items() if name in wanted_args}

    chunks = []
    for start in range(0, len(model_inputs), batch_size):
        batch = model_inputs[start : start + batch_size]
        encoded = tokenizer(
            batch,
            padding=True,
            truncation=True,
            max_length=max_length,
            return_tensors="pt",
        )
        input_ids = encoded["input_ids"].to(device)
        attention_mask = encoded.get("attention_mask", torch.ones_like(input_ids)).to(device)
        token_type_ids = torch.zeros_like(input_ids, dtype=torch.long, device=device)
        position_ids = torch.arange(input_ids.size(1), dtype=torch.long, device=device)

        kwargs = filter_kwargs(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
        )

        outputs = backbone(**kwargs)
        if hasattr(outputs, "last_hidden_state"):
            hidden = outputs.last_hidden_state
        elif isinstance(outputs, (tuple, list)) and len(outputs) > 0:
            hidden = outputs[0]
        else:
            hidden = outputs

        if not torch.is_tensor(hidden):
            raise TypeError(f"Backbone returned unsupported type: {type(hidden).__name__}")

        if use_cls_embed:
            embedding = hidden[:, 0, :]
        else:
            mask = attention_mask.unsqueeze(-1).type_as(hidden)
            embedding = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-6)

        chunks.append(embedding.cpu().numpy())

    return np.concatenate(chunks, axis=0)


def main():
    parser = argparse.ArgumentParser(description="Encode molecules with MolBERT and save embeddings")
    script_dir = Path(__file__).resolve().parent
    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("--model-class", default="molbert.models.smiles.SmilesMolbertModel")
    parser.add_argument("--ckpt-dir", default=str(script_dir / ".cache" / "molbert_100epochs"))
    parser.add_argument("--ckpt", default=None)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--max-length", type=int, default=128)
    parser.add_argument("--use-cls", action="store_true", default=False)
    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)

    model_inputs = [record["model_input"] for record in records]
    if not model_inputs:
        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)

    ckpt_path = Path(args.ckpt) if args.ckpt else Path(args.ckpt_dir) / "checkpoints" / "last.ckpt"
    if not ckpt_path.exists():
        print(f"Checkpoint not found: {ckpt_path}", file=sys.stderr)
        sys.exit(1)

    ckpt = safe_load_ckpt(str(ckpt_path))
    hp = ckpt.get("hyper_parameters") or ckpt.get("hparams") or {}
    namespace = hp_to_namespace(hp)

    model_class = import_obj(args.model_class)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model_class.load_from_checkpoint(
        checkpoint_path=str(ckpt_path),
        map_location=device,
        strict=False,
        args=namespace,
        weights_only=False,
    )

    tokenizer, reason = try_build_featurizer_adapter(namespace)
    if tokenizer is None:
        print(f"Could not attach tokenizer automatically: {reason}", file=sys.stderr)
        sys.exit(2)
    setattr(model, "tokenizer", tokenizer)

    embeddings = batch_embed_inputs(
        model_inputs,
        model,
        device,
        args.batch_size,
        args.max_length,
        args.use_cls,
    )

    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}.")
    if failures:
        print(f"Filtered {len(failures)} invalid rows. See {fail_path}.")


if __name__ == "__main__":
    main()
