import os
import torch
from uuid import uuid4
from pathlib import Path
from logging import Logger
from typing import List, Set, Optional

from transformers.models.auto.tokenization_auto import AutoTokenizer

from language_pipes.llm_layer_collector import LlmLayerCollector
from language_pipes.llm_layer_collector.auto.auto_rms import AutoRMSNorm
from language_pipes.llm_layer_collector.auto.auto_layer import AutoDecoderLayer
from language_pipes.llm_layer_collector.auto.static_auto_model import StaticAutoModel

from language_pipes.jobs.job import ComputeStep, Job
from language_pipes.jobs.job_data import computationStateToJobData

from language_pipes.modeling.llm_meta_data import LlmMetadata
from language_pipes.modeling.compute import compute_layers

from language_pipes.util import clone_model

class EndModel:
    model_id: str
    process_id: str
    device: str
    input_embedding: torch.nn.Embedding
    norm: AutoRMSNorm
    head: torch.nn.Linear
    collector: LlmLayerCollector
    layers: List[AutoDecoderLayer]

    def __init__(self, num_local_layers: int, model_dir: str, model_id: str, device: str, huggingface_token: str | None = None):
        self.model_id = model_id
        self.device = device

        self.process_id = str(uuid4())
        model_path = str(Path(model_dir) / self.model_id)
        if not os.path.exists(model_path):
            clone_model(model_id, model_path, token=huggingface_token)
        self.meta_data = LlmMetadata(model_path)
        self.collector = LlmLayerCollector(
            model_dir=os.path.join(model_path, 'data'),
            cache_file=os.path.join(model_path, 'cache.json'),
            device=device,
            dtype=torch.float16
        )
        self.layers = []
        if num_local_layers > 0:
            self.load_layers(num_local_layers)
        self.tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_path, 'data'))
    
    def load_layers(self, num_local_layers: int):
        self.layers = self.collector.load_layer_set(0, num_local_layers - 1, self.device)

    def compute_layers(self, job: Job):
        if job.data is None:
            raise Exception("Job did not have data")
        job.set_layer(
            state=compute_layers(0, job.data, self.device, self.layers, job.cache),
            layer=len(self.layers),
            num_hidden_layers=self.collector.config.num_hidden_layers
        )

    def size(self):
        return self.meta_data.embed_size + self.meta_data.head_size

    def load(self):
        self.input_embedding = self.collector.load_input_embedding(self.device)
        self.norm = self.collector.load_norm(self.device)
        self.head = self.collector.load_head(self.device)

    def tokenize(self, job: Job):
        prompt = self.tokenizer.apply_chat_template([m.to_json() for m in job.messages], tokenize=False, chat_template=self.tokenizer.chat_template, add_generation_prompt=True)
        input_tokens = [int(t) for t in self.tokenizer.encode(prompt, return_tensors='pt')[0].numpy()]
        job.input_ids = input_tokens
        job.prompt_tokens = len(input_tokens)
        job.next_step()

    def compute_embed(self, job: Job, logger: Logger, prefill_chunk_size: int):
        if job.compute_step != ComputeStep.EMBED and job.compute_step != ComputeStep.TOKENIZE:
            raise ValueError('Invalid step for embedding')
        if self.input_embedding is None:
            raise RuntimeError("Input Embedding must be loaded before computation")
        
        comp_state = StaticAutoModel.compute_embedding(
            prompt_tokens=job.prompt_tokens,
            chunk_size=prefill_chunk_size,
            input_embedder=self.input_embedding,
            input_ids=torch.tensor([job.input_ids]),
            config=self.collector.config,
            cache=job.cache
        )
        
        job.data = computationStateToJobData(comp_state)
        job.next_step()

    def compute_norm(self, job: Job):
        if job.data is None or job.data.state is None:
            raise RuntimeError("Cannot compute norm without job data")
        norm = self.norm(job.data.state.to(self.device))
        job.set_norm(norm)

    @staticmethod
    def _add_stop_tokens(stop_tokens: Set[int], token_value: Optional[int]):
        if token_value is None:
            return

        if isinstance(token_value, int) and token_value > 0:
            stop_tokens.add(token_value)
            return

        if isinstance(token_value, (list, tuple, set)):
            for token in token_value:
                if isinstance(token, int) and token > 0:
                    stop_tokens.add(token)

    def compute_head(self, job: Job):
        if self.head is None:
            raise RuntimeError("Head must be loaded before computation")
        if job.data is None or job.data.state is None:
            raise RuntimeError("Cannot compute head without job data")
        
        head = StaticAutoModel.compute_head(
            head=self.head, 
            state=job.data.state, 
            device=self.device,
            top_k=job.top_k,
            top_p=job.top_p,
            min_p=job.min_p,
            temperature=job.temperature
        )

        stop_tokens: Set[int] = set()

        EndModel._add_stop_tokens(stop_tokens, self.collector.config.eos_token_id)
        EndModel._add_stop_tokens(stop_tokens, self.tokenizer.eos_token_id)
        EndModel._add_stop_tokens(stop_tokens, self.tokenizer.convert_tokens_to_ids("<|eot_id|>"))

        job.set_output(head, stop_tokens)
        job.delta = self.tokenizer.decode([job.input_ids[-1]], skip_special_tokens=True)

    def set_result(self, job: Job):
        res_tokens = job.input_id_tensor()
        if res_tokens is None:
            raise Exception("Cannot decode result tensor: no input ids")
        job.result = self.tokenizer.decode(res_tokens[job.prompt_tokens:], skip_special_tokens=True)

    def clean_up(self):
        del self.input_embedding
        del self.norm
        del self.head
