import os
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import time

# Identity and Developer info
DEVELOPER = "Mik Tse Agency"
MODEL_NAME = "ChatPBC V4 Advanced" # Will be adjusted for V3.3
SYSTEM_PROMPT = f"You are ChatPBC, an expert AI business strategist developed by {DEVELOPER}. Provide strategic, actionable business advice."

# Configuration
MODEL_ID = os.environ.get("MODEL_ID", "chatpbc1/chatpbc-v4")
HF_TOKEN = os.environ.get("HF_TOKEN")

print(f"Loading model {MODEL_ID}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto",
    token=HF_TOKEN
)
tokenizer.pad_token = tokenizer.eos_token

def chat_function(message, history):
    # Handle file content if present (Gradio 4+ handles this in message dict)
    text = message["text"] if isinstance(message, dict) else message
    files = message["files"] if isinstance(message, dict) and "files" in message else []
    
    file_content = ""
    if files:
        for f in files:
            try:
                with open(f, "r", errors="ignore") as file:
                    file_content += f"\n[File: {os.path.basename(f)}]\n{file.read()[:2000]}\n"
            except: pass

    full_user_msg = text + (f"\n\nContext from uploaded files:\n{file_content}" if file_content else "")
    
    # Format Llama-2 chat prompt
    prompt = f"[INST] <<SYS>>\n{SYSTEM_PROMPT}\n<</SYS>>\n\n"
    for user_msg, assistant_msg in history:
        # history elements can be strings or dicts in newer Gradio
        u = user_msg["text"] if isinstance(user_msg, dict) else user_msg
        a = assistant_msg["text"] if isinstance(assistant_msg, dict) else assistant_msg
        prompt += f"{u} [/INST] {a} <s>[INST] "
    prompt += f"{full_user_msg} [/INST]"
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            do_sample=True,
            repetition_penalty=1.1,
            eos_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(output[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    return response

demo = gr.ChatInterface(
    fn=chat_function,
    title=MODEL_NAME,
    description=f"Expert AI Business Strategist developed by {DEVELOPER}",
    multimodal=True,
    theme="soft"
)

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