import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel
import os

ADAPTER_PATH = "KirubaLS/Outreach_model_6"
BASE_MODEL_NAME = "google/gemma-2b"
MAX_SEQ_LENGTH = 1024
MAX_NEW_TOKENS = 250

def load_model_and_tokenizer(adapter_path, base_model_name, hf_token=None):
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from peft import PeftModel

    device = torch.device("cpu")
    print(f"Using device: {device}")

    try:
        base_model = AutoModelForCausalLM.from_pretrained(
            base_model_name,
            torch_dtype=torch.float16,       
            low_cpu_mem_usage=True,          
            token=hf_token,         
            device_map="auto"
        )
        print("Base model loaded successfully.")
    except Exception as e:
        print(f"Failed to load base model: {e}")
        raise

    try:
        model = PeftModel.from_pretrained(
            base_model,
            adapter_path,
            token=hf_token
        )
        print("PEFT adapter loaded and attached successfully.")
    except Exception as e:
        print(f"Failed to load PEFT adapter from {adapter_path}: {e}")
        raise

    try:
        tokenizer = AutoTokenizer.from_pretrained(
            base_model_name,
            token=hf_token
        )
        print("Tokenizer loaded successfully.")
    except Exception as e:
        print(f"Failed to load tokenizer: {e}")
        raise

    model.eval()
    return model, tokenizer, device

try:
    model, tokenizer, device = load_model_and_tokenizer(ADAPTER_PATH, BASE_MODEL_NAME,hf_token=None)
except Exception as e:
    print(f"Error loading model and tokenizer: {e}")
    model = None
    tokenizer = None
    device = None


def generate_outreach_message(industry, company, prospect, designation, client_info, my_info):
    if model is None or tokenizer is None:
        return "Error: Model or tokenizer not loaded."

    input_prompt = (
        f"Industry: <industry>{industry}</industry>, "
        f"Company: <company>{company}</company>, "
        f"Prospect: <prospect>{prospect}</prospect>, "
        f"Designation: <designation>{designation}</designation>, "
        f"Client Info: <client_info>{client_info}</client_info>, "
        f"My Info: <my_info>{my_info}</my_info>, "
        f"Message Level: <message_level>First Level</message_level>: "
        f"{industry}\nHi {prospect.split(' ')[0]},\n"
    )

    truncated_input_prompt = input_prompt
    if len(tokenizer(input_prompt)['input_ids']) > MAX_SEQ_LENGTH:
        truncated_input_prompt = tokenizer.decode(tokenizer(input_prompt)['input_ids'][:MAX_SEQ_LENGTH])

    input_ids = tokenizer(
        truncated_input_prompt,
        return_tensors="pt",
        padding=False,
        truncation=True,
        max_length=MAX_SEQ_LENGTH
    ).to(device)

    with torch.no_grad():
        output_tokens = model.generate(
            input_ids["input_ids"],
            max_new_tokens=MAX_NEW_TOKENS,
            num_return_sequences=1,
            no_repeat_ngram_size=2,
            top_k=50,
            top_p=0.95,
            temperature=0.7,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    full_output_text = tokenizer.decode(output_tokens[0], skip_special_tokens=False)

    input_prompt_in_output_index = full_output_text.find(truncated_input_prompt)

    if input_prompt_in_output_index != -1:
        generated_text = full_output_text[input_prompt_in_output_index + len(truncated_input_prompt):].strip()
    else:
        input_token_length = input_ids["input_ids"].shape[1]
        generated_only_tokens = output_tokens[0, input_token_length:]
        generated_text = tokenizer.decode(generated_only_tokens, skip_special_tokens=True).strip()

    return generated_text

if __name__ == "__main__":
    if model is not None and tokenizer is not None:
        try:
            df = pd.read_csv('dataset.csv')
            data_row = df.loc[0]
            generated_message = generate_outreach_message(
                data_row['Industry'],
                data_row['Company Name'],
                data_row['Prospect Name'],
                data_row['Designation'],
                data_row['Client Info'],
                data_row['My Info']
            )
            print("\n--- Test Generation ---")
            print(generated_message)
            print("-----------------------")
        except FileNotFoundError:
            print("dataset.csv not found. Cannot run test generation.")
    else:
        print("Model and tokenizer not loaded, cannot run test generation.")