import os
import sys
import json
import unittest
import subprocess
from time import time
from pathlib import Path
from typing import Tuple, List, Dict

import torch
from torch import tensor
from transformers import AutoTokenizer
from transformers.cache_utils import DynamicCache

sys.path.append(os.path.join(os.path.dirname(__file__), '../../src'))

from language_pipes.llm_layer_collector.layer_collector import LlmLayerCollector
from language_pipes.llm_layer_collector.cache import get_shard_files
from language_pipes.llm_layer_collector.helpers import load_shard_tensor
from language_pipes.llm_layer_collector.load_layer import files_to_load_for_layer
from language_pipes.llm_layer_collector import StaticAutoModel
from language_pipes.util import clone_model

PROMPT = "The quick brown fox jumps over the "

with open('mcbeth.txt', 'r', encoding='utf-8') as f:
    mcbeth = f.read()

def get_base_dir(model_id: str):
    return str(Path(Path.home()) / ".cache" / "language_pipes" / "models" / model_id)

def get_model_dir(model_id: str):
    return get_base_dir(model_id) + "/data"

def get_cache_file(model_id: str):
    return get_base_dir(model_id) + "/cache.json"

def ensure_model(model_id: str):
    model_dir = get_base_dir(model_id)
    if os.path.exists(model_dir):
        return
    clone_model(model_id, model_dir)

def check_cache(tst: unittest.TestCase, model_dir: str, cache_file: str, num_keys: int):
    collector = LlmLayerCollector(model_dir, cache_file)
    tst.assertEqual(len(collector.layer_files.keys()), num_keys)
    tst.assertTrue(os.path.exists(cache_file))
    tst.assertTrue(os.path.exists(collector.cache_file))
    with open(cache_file, 'r') as f:
        cache = json.load(f)
        tst.assertEqual(len(cache.keys()), num_keys)

def get_tokenizer(model_dir: str) -> AutoTokenizer:
    return AutoTokenizer.from_pretrained(model_dir) # type: ignore

def encode_tokenizer(tokenizer: AutoTokenizer, prompt: str) -> torch.Tensor:
    return tokenizer(prompt, return_tensors='pt')['input_ids'] # type: ignore

def check_embedding(tst: unittest.TestCase, model_dir: str, cache_file: str, state_shape: Tuple[int, int, int], position_ids_shape: Tuple[int, int], position_embeddings_shape: Tuple[int, int, int]):
    collector = LlmLayerCollector(model_dir, cache_file)
    input_embedder = collector.load_input_embedding()
    tokenizer = get_tokenizer(model_dir)
    input_ids = encode_tokenizer(tokenizer, PROMPT)
    state = StaticAutoModel.compute_embedding(len(input_ids), 128, input_embedder, input_ids, collector.config, DynamicCache())
    tst.assertEqual(state.state.shape, state_shape)
    tst.assertEqual(state.position_ids.shape, position_ids_shape)
    # tst.assertEqual(state.position_embeddings[0].shape, position_embeddings_shape)
    # tst.assertEqual(state.position_embeddings[1].shape, position_embeddings_shape)
    tst.assertIsNotNone(state.cache_position)
    tst.assertIsNotNone(state.causal_mask)
    tst.assertGreater(len(state.causal_mask.keys()), 1)
    
def check_norm(tst: unittest.TestCase, model_dir: str, cache_file: str, norm_dim: int):
    collector = LlmLayerCollector(model_dir, cache_file)
    norm = collector.load_norm()
    norm = norm.to(dtype=torch.float16)
    tst.assertEqual(norm.cls.weight.shape, (norm_dim,))

def check_head(tst: unittest.TestCase, model_dir: str, cache_file: str, head_shape: Tuple[int, int]):
    collector = LlmLayerCollector(model_dir, cache_file)
    head = collector.load_head()
    tst.assertEqual(head.weight.shape, head_shape)

def check_layers(tst: unittest.TestCase, model_dir: str, cache_file: str, end_layer: int):
    start_time = time()
    collector = LlmLayerCollector(model_dir, cache_file)
    layers = collector.load_layer_set(0, end_layer)
    print(f"Time: {time() - start_time:.2f}s")
    tst.assertEqual(len(layers), end_layer+1)

def apply_chat_template(tokenizer: AutoTokenizer, messages: List[Dict[str, str]]) -> torch.Tensor:
    return tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors='pt') # type: ignore

def decode_tokenizer(tokenizer: AutoTokenizer, ids: torch.Tensor) -> str:
    return tokenizer.decode(ids) # type: ignore

def check_stack(tst: unittest.TestCase, model_dir: str, cache_file: str, chunk_size: int = 8):
    tokenizer = get_tokenizer(model_dir)
    chat = [
        {"role": "system", "content": "You are a helpful assistant",},
        {"role": "user", "content": f"What play is this text from?\n{mcbeth[:500]}"},
    ]
    all_input_ids = apply_chat_template(tokenizer, chat)
    prompt_tokens = all_input_ids.shape[1]
    print(f"Total tokens: {prompt_tokens}, Chunk size: {chunk_size}")
    
    collector = LlmLayerCollector(model_dir, cache_file)
    input_embed = collector.load_input_embedding()
    head = collector.load_head()
    norm = collector.load_norm()
    layers = collector.load_layer_set(0, collector.config.num_hidden_layers - 1)
    
    cache = DynamicCache()
    
    num_chunks = (prompt_tokens + chunk_size - 1) // chunk_size
    print(f"Processing {num_chunks} chunks")

    for chunk_idx in range(num_chunks):
        state = StaticAutoModel.compute_embedding(prompt_tokens, chunk_size, input_embed, all_input_ids, collector.config, cache)
        
        tst.assertEqual(state.cache_position[0].item(), chunk_idx * chunk_size)
        tst.assertEqual(state.cache_position[-1].item(), min(((chunk_idx + 1) * chunk_size) - 1, prompt_tokens - 1))
        
        for lyr in layers:
            state.state = StaticAutoModel.compute_layer(lyr, state, cache)
        
        tst.assertEqual(cache.get_seq_length(), min((chunk_idx + 1) * chunk_size, prompt_tokens))
        print(f"chunk {chunk_idx + 1} / {num_chunks}")
    
    tst.assertEqual(cache.get_seq_length(), prompt_tokens)
    
    num_decode_tokens = 30
    current_ids = all_input_ids
    
    for i in range(num_decode_tokens):
        state = StaticAutoModel.compute_embedding(prompt_tokens, chunk_size, input_embed, current_ids, collector.config, cache)
        
        tst.assertEqual(state.state.shape[1], 1)
        
        for lyr in layers:
            state.state = StaticAutoModel.compute_layer(lyr, state, cache)
        
        result = StaticAutoModel.compute_head(head, norm(state.state), 'cpu', 10)
        
        token_list: List[int] = current_ids.cpu().numpy().tolist()[0]
        token_list.append(result)
        current_ids = tensor([token_list])
        
        print(f"Decode token {i + 1}: cache size = {cache.get_seq_length()}")
    
    print(f"Generated {num_decode_tokens} tokens after chunked prefill")
    print(f"Output: {decode_tokenizer(tokenizer, current_ids[0][prompt_tokens:])}")
    
    tst.assertEqual(current_ids.shape[1], prompt_tokens + num_decode_tokens)

class TestLlmLayerCollector(unittest.TestCase):
    @unittest.skip("")
    def test_qwen3_2B(self):
        model_id = "Qwen/Qwen3-1.7B"
        model_dir = get_model_dir(model_id)
        cache_file = get_cache_file(model_id)
        ensure_model(model_id)
        check_cache(self, model_dir, cache_file, 311)
        check_embedding(self, model_dir, cache_file, (1, 8, 2048), (1, 8), (1, 8, 128))
        check_norm(self, model_dir, cache_file, 2048)
        check_head(self, model_dir, cache_file, (151936, 2048))
        check_layers(self, model_dir, cache_file, 10)
        check_stack(self, model_dir, cache_file, chunk_size=32)

    @unittest.skip("")
    def test_llama3_8B(self):
        model_id = "meta-llama/Llama-3.1-8B-Instruct"
        model_dir = get_model_dir(model_id)
        cache_file = get_cache_file(model_id)
        ensure_model(model_id)
        check_cache(self, model_dir, cache_file, 291)
        check_embedding(self, model_dir, cache_file, (1, 9, 4096), (1, 9), (1, 9, 128))
        check_norm(self, model_dir, cache_file, 4096)
        check_head(self, model_dir, cache_file, (128256, 4096))
        check_layers(self, model_dir, cache_file, 2)
        check_stack(self, model_dir, cache_file, chunk_size=32)

    @unittest.skip("")
    def test_phi4_4B(self):
        model_id = "microsoft/Phi-4-mini-reasoning"
        model_dir = get_model_dir(model_id)
        cache_file = get_cache_file(model_id)
        ensure_model(model_id)
        check_cache(self, model_dir, cache_file, 194)
        check_embedding(self, model_dir, cache_file, (1, 8, 3072), (1, 8), (1, 8, 128))
        check_norm(self, model_dir, cache_file, 3072)
        check_head(self, model_dir, cache_file, (200064, 3072))
        check_layers(self, model_dir, cache_file, 2)
        check_stack(self, model_dir, cache_file, chunk_size=32)

    def test_gemma3_1B(self):
        model_id = "google/gemma-3-1b-it"
        model_dir = get_model_dir(model_id)
        cache_file = get_cache_file(model_id)
        ensure_model(model_id)
        check_cache(self, model_dir, cache_file, 340)
        check_embedding(self, model_dir, cache_file, (1, 9, 1152), (1, 9), (1, 9, 256))
        check_norm(self, model_dir, cache_file, 1152)
        check_head(self, model_dir, cache_file, (262144, 1152))
        check_layers(self, model_dir, cache_file, 2)
        check_stack(self, model_dir, cache_file, chunk_size=32)

    def test_glm41V_9B(self):
        model_id = "zai-org/GLM-4.1V-9B-Thinking"
        model_dir = get_model_dir(model_id)
        cache_file = get_cache_file(model_id)
        ensure_model(model_id)
        check_cache(self, model_dir, cache_file, 704)
        check_embedding(self, model_dir, cache_file, (1, 8, 4096), (3, 1, 8), (1, 8, 256)) # pyright: ignore[reportArgumentType]
        check_norm(self, model_dir, cache_file, 4096)
        check_head(self, model_dir, cache_file, (151552, 4096))
        check_layers(self, model_dir, cache_file, 2)
        check_stack(self, model_dir, cache_file, chunk_size=32)

    def test_exceptions(self):
        model_id = "Qwen/Qwen3-1.7B"
        model_dir = get_model_dir(model_id)
        cache_file = get_cache_file(model_id)
        collector = LlmLayerCollector(model_dir, cache_file)

        try:
            os.mkdir('shard_test')
            get_shard_files(collector.shard_pattern, 'shard_test')
            self.fail("Should have thrown an exception")
        except Exception:
            pass
        os.rmdir('shard_test')
        
        try:
            load_shard_tensor(collector.layer_files, collector.model_dir, 'bad_layer', torch.device('cpu'), torch.float16)
            self.fail("Should have thrown an exception")
        except ValueError:
            pass

        try:
            os.mkdir('bad_dir')
            LlmLayerCollector('bad_dir', cache_file)
            self.fail("Should have thrown an exception")
        except FileNotFoundError:
            pass

        os.rmdir('bad_dir')

        try:
            os.remove(cache_file)
            collector._read_cache() # type: ignore
            self.fail("Should have thrown an exception")
        except FileNotFoundError:
            pass
        
        try:
            files_to_load_for_layer('bad_key', []) # type: ignore
            self.fail("Should have thrown an exception")
        except Exception:
            pass

# If you want to run these tests directly from the command line:
if __name__ == '__main__':
    unittest.main()
