import torch
import torch.nn.functional as F
import os
import gradio as gr
from transformers import AutoTokenizer, AutoModel, AutoConfig

# Repository ID
MODEL_ID = "mahaveerai/bol-ai"
TEMPERATURE = 0.1
MAX_NEW_TOKENS = 300

print("Initializing Bol-AI v1.0 Pro Engine...")

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

# Handle custom model_type 'bol_ai_v1'
try:
    model = AutoModel.from_pretrained(
        MODEL_ID,
        dtype=torch.bfloat16,  # 'torch_dtype' warning fixed here
        device_map="auto",
        trust_remote_code=True
    )
except KeyError:
    print("Applying architecture mapping...")
    config = AutoConfig.from_pretrained("openbmb/MiniCPM-V-4.6", trust_remote_code=True)
    model = AutoModel.from_pretrained(
        MODEL_ID,
        config=config,
        dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True
    )

model.eval()

def custom_generate(user_input):
    """Manual generation loop to bypass missing .chat() or .generate() methods"""
    prompt = f"User: {user_input}\nBol-AI:"
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
    
    # Dynamically find the Language Model Head
    lm_head = None
    with torch.no_grad():
        out = model(input_ids)
        h = out.last_hidden_state if hasattr(out, "last_hidden_state") else out
        dim = h.shape[-1]
        
        for name, module in model.named_modules():
            if isinstance(module, torch.nn.Linear) and module.in_features == dim and module.out_features > 20000:
                lm_head = lambda x: module(x.to(module.weight.dtype))
                break
        if not lm_head:
            for module in model.modules():
                if isinstance(module, torch.nn.Embedding) and module.embedding_dim == dim and module.num_embeddings > 20000:
                    lm_head = lambda x: torch.matmul(x.to(module.weight.dtype), module.weight.T)
                    break
                    
    if not lm_head:
        return "Error: Language head not found."

    generated_ids = input_ids[0].tolist()  # Fixed slicing error here
    start_len = len(generated_ids)
    
    # Generate tokens
    for _ in range(MAX_NEW_TOKENS):
        curr_tensor = torch.tensor([generated_ids]).to(model.device)
        with torch.no_grad():
            out = model(curr_tensor)
            h = out.last_hidden_state if hasattr(out, "last_hidden_state") else out
            logits = lm_head(h[:, -1, :])
            
            token = torch.argmax(logits, dim=-1).item()
            generated_ids.append(token)
            
            if token == tokenizer.eos_token_id:
                break
                
    response = tokenizer.decode(generated_ids[start_len:], skip_special_tokens=True)
    return response.split("User:")[0].strip()  # Fixed string parsing logic here

# Gradio Interface Chat Setup
def chat_interface(message, history):
    return custom_generate(message)

# Launching Web UI with explicit theme block to fix TypeError
with gr.Blocks(theme="soft") as demo:
    gr.ChatInterface(
        fn=chat_interface,
        title="BOL-AI v1.0 PRO",
        description="Developer: Vivek Vijay Dalvi | Company: MAHAVEER AI"
    )

if __name__ == "__main__":
    demo.launch()