import gradio as gr
import torch
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from langchain_classic.chains import create_retrieval_chain
from langchain_classic.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import PromptTemplate

# 1. Vector Store Setup
embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
vectorstore = FAISS.load_local(
    "faiss_upf_index", 
    embeddings, 
    allow_dangerous_deserialization=True
)
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# 2. Model & Pipeline Initialization
model_id = "anirudh248/upf-code-generator"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto"
)
model.generation_config.pad_token_id = tokenizer.eos_token_id

hf_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=1024,
    temperature=0.2,
    top_p=0.95,
    repetition_penalty=1.15,
    return_full_text=False,
    clean_up_tokenization_spaces=False
)
llm = HuggingFacePipeline(pipeline=hf_pipeline)

# 3. RAG Chain Setup
unified_prompt_template = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>

You are a highly capable AI assistant specializing in Unified Power Format (UPF 3.0) and VLSI power intent design.

Instructions:
1. If the User Input asks for UPF code or power intent, act as an expert UPF engineer. Use the Context to generate precise UPF 3.0 code, enclosed in ```tcl ... ``` blocks.
2. If the User Input is a general question or greeting, respond conversationally and concisely. Ignore the Context if it is not relevant.

Context:
{context}<|eot_id|><|start_header_id|>user<|end_header_id|>

Conversation History:
{chat_history}

User Input: {input}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
"""

unified_prompt = PromptTemplate.from_template(unified_prompt_template)
document_chain = create_stuff_documents_chain(llm, unified_prompt)
rag_chain = create_retrieval_chain(retriever, document_chain)

# 4. Gradio Interface & Handling
def get_text_content(msg_content):
    """Safety wrapper to ensure Gradio doesn't accidentally pass JSON dicts to the LLM"""
    if isinstance(msg_content, str):
        return msg_content
    elif isinstance(msg_content, list):
        return " ".join([item.get('text', '') for item in msg_content if isinstance(item, dict) and 'text' in item])
    return str(msg_content)

def format_history(history):
    if not history:
        return "No previous conversation."
    return "\n".join([f"{msg['role'].capitalize()}: {get_text_content(msg['content'])}" for msg in history])

def user_interaction(user_message, history):
    history = history or []
    user_text = get_text_content(user_message)
    
    response = rag_chain.invoke({
        "input": user_text,
        "chat_history": format_history(history)
    })
    
    # Strip any trailing end-of-turn tokens Llama 3 might accidentally output
    answer = response['answer'].replace("<|eot_id|>", "").strip()
    
    upf_syntax_hints = ["create_power_domain", "set_isolation", "set_retention", "create_supply_port", "create_power_switch"]
    if "```" not in answer and any(kw in answer.lower() for kw in upf_syntax_hints):
        answer = f"```tcl\n{answer}\n```"

    history.append({"role": "user", "content": user_text})
    history.append({"role": "assistant", "content": answer})
    return history, ""

with gr.Blocks() as interface:
    gr.Markdown("# UPF Code Generator with Llama 3 & RAG")
    
    chatbot = gr.Chatbot(label="Chat History", elem_id="chatbot")
    
    with gr.Row():
        user_input = gr.Textbox(show_label=False, placeholder="Enter a general question or a UPF power intent...", lines=3)
        
    send_button = gr.Button("Generate Response", variant="primary")
    
    send_button.click(fn=user_interaction, inputs=[user_input, chatbot], outputs=[chatbot, user_input])
    user_input.submit(fn=user_interaction, inputs=[user_input, chatbot], outputs=[chatbot, user_input])

interface.launch(theme=gr.themes.Soft())