import torch
from PIL import Image
import requests
from dataclasses import dataclass, field
from transformers.hf_argparser import HfArgumentParser
from transformers import (
    AutoConfig,
    AutoImageProcessor,
    AutoModel,
    AutoModelForCausalLM,
    AutoTokenizer,
    LlavaNextConfig,
    LlavaNextForConditionalGeneration,
    LlavaNextImageProcessor,
    LlavaNextProcessor,
)
from src.logger import get_logger


logger = get_logger(__name__)


@dataclass
class CreateModelArguments:
    vision_tower_model_name_or_path: str = field(
        default="openai/clip-vit-large-patch14-336",
        metadata={"help": "Path to the vision tower model directory."},
    )
    text_model_name_or_path: str = field(
        default="CYFRAGOVPL/Llama-PLLuM-8B-instruct",
        metadata={"help": "Path to the text model directory."},
    )
    save_model_path: str = field(
        default="models/Llama-PLLuM-8B-instruct-llava-next",
        metadata={"help": "Path to save the created model."},
    )
    check_model: bool = field(
        default=False,
        metadata={"help": "Whether to check if the model can be loaded and used after saving."},
    )


def create_model(
    args: CreateModelArguments,
) -> tuple[LlavaNextForConditionalGeneration, AutoTokenizer, LlavaNextProcessor]:
    logger.info("Creating the model...")

    vit_model = AutoModel.from_pretrained(args.vision_tower_model_name_or_path)
    vit_config = AutoConfig.from_pretrained(args.vision_tower_model_name_or_path)
    vision_config = getattr(vit_config, "vision_config", vit_config)
    logger.info("Loaded the vision tower model from " + args.vision_tower_model_name_or_path)

    text_model = AutoModelForCausalLM.from_pretrained(
        args.text_model_name_or_path,
        dtype=torch.bfloat16,
    )
    text_config = AutoConfig.from_pretrained(args.text_model_name_or_path)
    logger.info("Loaded the text model from " + args.text_model_name_or_path)

    tokenizer = AutoTokenizer.from_pretrained(args.text_model_name_or_path)
    tokenizer.add_special_tokens({"additional_special_tokens": ["<image>"]})
    image_token_id = tokenizer.convert_tokens_to_ids("<image>")
    logger.info("Loaded the tokenizer and added special tokens")

    logger.info("Creating the LlavaNextForConditionalGeneration model...")
    config = LlavaNextConfig(
        text_config=text_config,
        vision_config=vision_config,
        image_token_index=image_token_id,
        dtype=torch.bfloat16,
        vision_feature_select_strategy="full",
        image_seq_length=(vision_config.image_size // vision_config.patch_size) ** 2,
        image_grid_pinpoints=[[384, 768], [768, 384], [768, 768], [1152, 384], [384, 1152]],
        use_image_newline_parameter=True,
    )

    model = LlavaNextForConditionalGeneration(config=config)
    logger.info("Created the LLaVA-Next model")

    model.language_model.load_state_dict(text_model.model.state_dict())
    model.lm_head.load_state_dict(text_model.lm_head.state_dict())
    logger.info("Loaded text model weights")

    if hasattr(vit_model, "vision_model"):
        vision_state = vit_model.vision_model.state_dict()
    else:
        vision_state = vit_model.state_dict()

    model.vision_tower.vision_model.load_state_dict(vision_state, strict=True)
    logger.info("Loaded vision tower weights")

    model.resize_token_embeddings(len(tokenizer))
    logger.info("Resized embeddings to match the tokenizer size and saved the model")

    vit_image_processor = AutoImageProcessor.from_pretrained(args.vision_tower_model_name_or_path)
    target = vision_config.image_size
    image_processor = LlavaNextImageProcessor(
        size={"shortest_edge": target},
        resample=3,
        crop_size={"height": target, "width": target},
        image_mean=vit_image_processor.image_mean,
        image_std=vit_image_processor.image_std,
        image_grid_pinpoints=config.image_grid_pinpoints,
    )

    processor = LlavaNextProcessor(
        image_processor=image_processor,
        tokenizer=tokenizer,
        patch_size=vision_config.patch_size,
    )
    processor.num_additional_image_tokens = 0
    processor.vision_feature_select_strategy = "full"
    logger.info("Created the LlavaNextProcessor")

    return model, tokenizer, processor


def save_model(model, tokenizer, processor, args: CreateModelArguments):
    model.save_pretrained(args.save_model_path)
    tokenizer.save_pretrained(args.save_model_path)
    processor.save_pretrained(args.save_model_path)
    logger.info("Saved the model, tokenizer and processor to " + args.save_model_path)


def check_model(args: CreateModelArguments):
    logger.info("Checking if the model can be loaded and used from the saved path...")
    try:
        model = LlavaNextForConditionalGeneration.from_pretrained(
            args.save_model_path,
            dtype=torch.bfloat16,
            device_map="auto",
            low_cpu_mem_usage=True,
        )

        processor = LlavaNextProcessor.from_pretrained(args.save_model_path)

        conversation = [{"role": "user", "content": "<image>\nCo jest pokazane na obrazku?"}]
        prompt = processor.apply_chat_template(
            conversation, tokenize=False, add_generation_prompt=True
        )
        logger.info(f"Prompt for the model: {prompt}")

        image_url = "https://www.ilankelman.org/stopsigns/australia.jpg"
        image = Image.open(requests.get(image_url, stream=True).raw)
        logger.info(f"Loaded image from {image_url}")

        inputs = processor(
            images=image, text=prompt, return_tensors="pt", return_token_type_ids=False
        )
        inputs = inputs.to(model.device)

        generate_ids = model.generate(**inputs, max_new_tokens=64, do_sample=False)
        generate_text = processor.batch_decode(
            generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
        )[0]

        logger.info(f"Model loaded successfully. Example generated output: {generate_text}")
    except Exception as e:
        logger.error(f"Failed to load model from {args.save_model_path}: {e}")


if __name__ == "__main__":
    parser = HfArgumentParser(CreateModelArguments)
    (args,) = parser.parse_args_into_dataclasses()

    model, tokenizer, processor = create_model(args)
    save_model(model, tokenizer, processor, args)

    if args.check_model:
        check_model(args)

    logger.info("Model creation and saving completed successfully.")
