import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_huggingface import HuggingFacePipeline
import os

# --- 1. Model ko Hugging Face se load karein ---
# (Space shuru hone par yeh model automatically download aur load karega)
print("Model ko load kiya ja raha hai...")
model_id = "google/gemma-2b-it"
quantization_config = BitsAndBytesConfig(load_in_4bit=True)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=quantization_config,
    device_map="auto"
)
print("Model safaltapoorvak load ho gaya!")

# --- 2. LLM Pipeline Banayein ---
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=1024)
llm = HuggingFacePipeline(pipeline=pipe)

# --- 3. Agent ki sabhi kshamtaein (Chains) Banayein ---
# (Yahan hum apne sabhi purane prompt templates daalenge)

# Text-to-Code
text_to_code_chain = LLMChain(llm=llm, prompt=PromptTemplate(input_variables=["instruction"], template="You are an expert Python coder. Write clean Python code for this instruction:\n\n{instruction}"))

# Code-to-Text
code_to_text_chain = LLMChain(llm=llm, prompt=PromptTemplate(input_variables=["code"], template="You are an expert code reviewer. Explain this code in simple points:\n\nCODE:\n{code}\n\nEXPLANATION:"))

# Code Translation
code_translation_chain = LLMChain(llm=llm, prompt=PromptTemplate(input_variables=["source_language", "target_language", "code"], template="Translate this {source_language} code to {target_language}:\n\n{code}"))

# Code Refactoring
code_refactoring_chain = LLMChain(llm=llm, prompt=PromptTemplate(input_variables=["code"], template="Refactor this code to be more efficient and readable:\n\n{code}"))

# Debugging
debugging_chain = LLMChain(llm=llm, prompt=PromptTemplate(input_variables=["problem_description", "code"], template="Debug this code. Problem: {problem_description}\n\nCODE:\n{code}"))


# --- 4. Agent ka Logic Function ---
def run_agent(task, user_input):
    print(f"Agent '{task}' karya chala raha hai...")
    if task == "text_to_code":
        return text_to_code_chain.run(user_input)
    elif task == "code_to_text":
        return code_to_text_chain.run(user_input)
    elif task == "code_translation":
        return code_translation_chain.run(user_input)
    elif task == "code_refactoring":
        return code_refactoring_chain.run(user_input)
    elif task == "debugging":
        return debugging_chain.run(user_input)
    else:
        return "Anya karya. Kripya uplabdh karyon mein se chunein."

# --- 5. Gradio ka UI Function ---
def agent_ui(task, user_input):
    if not task: return "⚠️ Kripya ek karya chunein."
    if not user_input.strip(): return "⚠️ Kripya input text box mein kuch likhein."
    try:
        if task == "code_translation":
            parts = user_input.split('|||')
            if len(parts) != 3: return "Error: Translation ke liye format: source_lang ||| target_lang ||| code"
            input_dict = {"source_language": parts[0].strip(), "target_language": parts[1].strip(), "code": parts[2].strip()}
            return run_agent(task, input_dict)
        elif task == "debugging":
            parts = user_input.split('|||')
            if len(parts) != 2: return "Error: Debugging ke liye format: problem_description ||| code"
            input_dict = {"problem_description": parts[0].strip(), "code": parts[2].strip()}
            return run_agent(task, input_dict)
        else:
            return run_agent(task, user_input)
    except Exception as e:
        return f"Ek error hui: {e}"

# --- 6. Gradio Interface Banayein ---
task_dropdown = gr.Dropdown(choices=["text_to_code", "code_to_text", "code_translation", "code_refactoring", "debugging"], label="Karya Chunein (Select Task)")
input_textbox = gr.Textbox(lines=10, label="Aapka Input (Your Input)", placeholder="Yahan apna nirdesh ya code paste karein...")
output_markdown = gr.Markdown(label="Agent ka Jawab (Agent's Response)")

iface = gr.Interface(
    fn=agent_ui,
    inputs=[task_dropdown, input_textbox],
    outputs=output_markdown,
    title="🤖 AI Coding Agent",
    description="Namaste! Main aapka AI coding assistant hoon. Ek karya chunein aur neeche apna nirdesh likhein."
)

# --- 7. App ko Launch Karein ---
iface.launch()