import os
import streamlit as st
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from langchain.prompts import FewShotPromptTemplate, PromptTemplate
from langchain.chains import LLMChain
from langchain_groq import ChatGroq
import uvicorn

# Load environment variables
load_dotenv()

# Set up API key
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

if not GROQ_API_KEY:
    raise ValueError("🚨 API Key Missing! Please check your .env file and restart the app.")

llm = ChatGroq(api_key=GROQ_API_KEY, model_name="gemma2-9b-it")

# Define example responses for few-shot prompting
examples = [
    {
        "input": "def add(a, b):\nreturn a + b",
        "output": "Your function 'add' is missing proper indentation. Here's a corrected version:\n\ndef add(a, b):\n    return a + b\n"
    },
    {
        "input": "def divide(a, b):\n    return a / b",
        "output": "Potential bug detected: Division by zero error. You should handle this case:\n\ndef divide(a, b):\n    if b == 0:\n        return 'Error: Division by zero'\n    return a / b\n"
    }
]

# Define example template
example_template = PromptTemplate(
    input_variables=["input", "output"],
    template="Code: \n{input}\n\nFeedback:\n{output}\n"
)
prefix="""You are a highly skilled Python code reviewer. 
Your task is to analyze the given Python code, identify potential bugs, suggest improvements, and provide a corrected version of the code if necessary. Ensure that your feedback is clear, precise, and actionable.
First you have to specify where and what the error is.
Next give the correct code
If the code is out of context replay "Out of Context"

"""

# Create a few-shot prompt template
few_shot_prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=example_template,
    prefix=prefix,
    suffix="Code:\n{input}\n\nFeedback:",
    input_variables=["input"]
)

# Create the LLMChain
llm_chain = LLMChain(llm=llm, prompt=few_shot_prompt)

# FastAPI Backend
app = FastAPI()

class CodeReviewRequest(BaseModel):
    code: str

@app.post("/review")
def review_code(request: CodeReviewRequest):
    if not request.code.strip():
        raise HTTPException(status_code=400, detail="No code provided.")
    response = llm_chain.run(input=request.code)
    return {"feedback": response}

# Streamlit Frontend
st.title("🤖 AI Code Reviewer 📝")
st.markdown("### Get instant feedback on your Python code! 🚀")

code_snippet = st.text_area("✍️ Paste your Python code below:", height=200)

if st.button("🔍 Review Code"):
    if code_snippet.strip():
        response = llm_chain.run(input=code_snippet)
        st.subheader("🧐 Review Feedback:")
        st.code(response, language="python")
    else:
        st.warning("⚠️ Please enter some Python code to review!")

# Run FastAPI backend (for local testing)
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8001)