import gradio as gr
import torch
import os
import json
import uuid
from datetime import datetime
from pathlib import Path
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
from huggingface_hub import CommitScheduler

# ==========================================
# 1. CONFIGURATION & LOGGING SETUP
# ==========================================
MODEL_ID = "sensix-zo/sensix-paite-4b-multiturn-16bit"
HF_TOKEN = os.getenv("HF_TOKEN")

# --- MATCHING YOUR SCREENSHOT ---
CHAT_LOG_DATASET_ID = "sensix-zo/sensix-chat-logs" 

log_folder = Path("logs")
log_folder.mkdir(parents=True, exist_ok=True)

# Syncs the 'logs' folder to HF 'data/' every 1 MINUTE
scheduler = CommitScheduler(
    repo_id=CHAT_LOG_DATASET_ID,
    repo_type="dataset",
    folder_path=log_folder,
    path_in_repo="data", 
    token=HF_TOKEN,
    every=1 
)

def save_to_logs(full_history):
    """Saves the full session history as a single JSON line in the master file."""
    log_file_path = log_folder / "chat_dataset.jsonl"
    try:
        with scheduler.lock:
            with log_file_path.open("a", encoding="utf-8") as f:
                entry = {"messages": full_history}
                f.write(json.dumps(entry, ensure_ascii=False) + "\n")
    except Exception as e:
        print(f"Logging error: {e}")

# ==========================================
# 2. LOAD TOKENIZER & MODEL (CPU)
# ==========================================
print("🚀 System: Initializing Gemma-3 on CPU...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)

device = "cpu"
dtype = torch.bfloat16 

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=dtype,
    device_map="cpu",
    token=HF_TOKEN,
    low_cpu_mem_usage=True
)

# ==========================================
# 3. CHAT HANDLER 
# ==========================================
def chat_handler(message, history):
    
    current_history = []
    
    # Safely handle the history regardless of Gradio version
    for item in history:
        if isinstance(item, dict):
            current_history.append({"role": item["role"], "content": item["content"]})
        elif isinstance(item, (list, tuple)) and len(item) == 2:
            user_msg, bot_msg = item
            if user_msg: current_history.append({"role": "user", "content": user_msg})
            if bot_msg: current_history.append({"role": "assistant", "content": bot_msg})

    # Add the current user prompt
    user_text = message if isinstance(message, str) else message.get("content", "")
    current_history.append({"role": "user", "content": user_text})

    # Prepare inputs for the model
    model_inputs = tokenizer.apply_chat_template(
        [{"role": m["role"].replace("assistant", "model"), "content": m["content"]} for m in current_history],
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True 
    ).to(device)

    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    generate_kwargs = dict(
        **model_inputs, 
        streamer=streamer,
        max_new_tokens=1024,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id else tokenizer.eos_token_id
    )

    with torch.no_grad():
        thread = Thread(target=model.generate, kwargs=generate_kwargs)
        thread.start()
        
        response_text = ""
        for new_token in streamer:
            response_text += new_token
            yield response_text
        
        # Build the final multiturn sequence for this turn
        final_history = current_history.copy()
        final_history.append({"role": "assistant", "content": response_text})
        
        # Append the complete history of this session as one JSON line to the master file
        save_to_logs(final_history)

# ==========================================
# 4. UI SETUP
# ==========================================
demo = gr.ChatInterface(
    fn=chat_handler,
    title="Sensix Paite AI",
    description="Please note: This AI is currently running on standard CPU hardware, so responses may take a little longer than usual.", 
    examples=["Na dam maw?", "Java bang a?", "How are you? chih Paite pau in bang chi gen di?"],
    cache_examples=False,
)

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