import streamlit as st
import os
import json
from openai import OpenAI
from dotenv import load_dotenv
import re

load_dotenv()

client = OpenAI(
  base_url="https://openrouter.ai/api/v1",
  api_key=os.getenv("OPENROUTER_API_KEY"),
)

def generate_code_analysis_with_retry(code_snippet: str, model_name: str,language: str, max_retries: int = 3):
    """
    Generates code analysis with a self-correction loop.
    It tries to get valid JSON, and if it fails, it tells the AI its mistake and retries.
    Returns a parsed dictionary on success, or None on failure.
    """
    
    # Define the initial user request
    initial_prompt = f"""
    You are an expert {language} programmer. Analyze the following code snippet and provide a plain-English explanation and a Google-style docstring.
    
    Code:
    ```
    {code_snippet}
    ```
    
    Respond with ONLY a single, valid JSON object with two keys: "explanation" and "docstring". Do not include any markdown formatting, comments, or other text outside of the JSON.
    """
    
    # Initialize the conversation history for the AI
    messages = [{"role": "user", "content": initial_prompt}]
    
    # Start the self-correction loop
    for attempt in range(max_retries):
        st.write(f" Attempt {attempt + 1} of {max_retries}...") 
        
        try:
            # === ACT: Call the AI ===
            response = client.chat.completions.create(
                model=model_name,
                messages=messages
            )
            raw_output = response.choices[0].message.content
            
            
            # Use regex to find the JSON object within the potentially messy string
            match = re.search(r"\{.*\}", raw_output, re.DOTALL)
            if not match:
                raise ValueError("No JSON object found in the response.")
            
            cleaned_json_str = match.group(0)
            
            parsed_json = json.loads(cleaned_json_str)
            
            if "explanation" not in parsed_json or "docstring" not in parsed_json:
                raise ValueError("JSON is missing required keys ('explanation', 'docstring').")

            st.success(f"Analysis successful on attempt {attempt + 1}!")
            return parsed_json 

        except (json.JSONDecodeError, ValueError, IndexError) as e:
            # === REASON & REACT: If an error occurred, start the correction process ===
            st.warning(f"Attempt {attempt + 1} failed: {e}. Trying to self-correct...")
            
            # Add the AI's failed response to the conversation history
            messages.append({"role": "assistant", "content": raw_output})
            
            # Create the corrective prompt, showing the AI its own mistake
            corrective_prompt = f"""
            Your previous response could not be parsed.
            Error: "{e}"
            Your full response was:
            ---
            {raw_output}
            ---
            Please correct your mistake. Look at the error and your previous response. 
            Provide the response again as a single, valid JSON object with the keys "explanation" and "docstring". 
            DO NOT wrap it in markdown or add any other text.
            """
            
            # Add  corrective instruction to the conversation
            messages.append({"role": "user", "content": corrective_prompt})
            
    st.error(f"Failed to get a valid response after {max_retries} attempts.")
    return None



st.set_page_config(layout="wide")
st.title("AI Code Explainer & Docstring Generator")
st.write("Powered by OpenRouter.ai with a Self-Correction Loop")

code_input = st.text_area(
    "Paste your Python function or code block here:", 
    height=250, 
    placeholder="def my_function(arg1, arg2):\n    # Your code here\n    return result"
)

model_choice = st.selectbox(
    "Choose your AI model:",
    (
        "Google: Gemma 3n",
        "MoonshotAI: Kimi Dev ",
        "NVIDIA: Nemotron Nano 9B",
        "Mistral: Mistral 7B Instruct",
    ),
    help="Free models from OpenRouter. Different models have different strengths."
)

MODEL_MAPPING = {
    "Google: Gemma 3n": "google/gemma-3n-e2b-it:free",
    "MoonshotAI: Kimi Dev ": "moonshotai/kimi-dev-72b:free",
    "NVIDIA: Nemotron Nano 9B": "nvidia/nemotron-nano-9b-v2:free",
    "Mistral: Mistral 7B Instruct": "mistralai/mistral-7b-instruct:free",
}
selected_model_id = MODEL_MAPPING[model_choice]

language = st.selectbox("Select Language", ["Python", "JavaScript", "Java", "Go"])
if st.button("Analyze Code", type="primary"):
    if code_input:
        analysis_dict = generate_code_analysis_with_retry(code_input, selected_model_id,language)
        
        if analysis_dict:
            st.subheader("Final Analysis Results")
            col1, col2 = st.columns(2)
            
            with col1:
                st.info("💬 Plain English Explanation")
                st.write(analysis_dict.get("explanation", "No explanation was generated."))
            
            with col2:
                st.success("📝 Generated Docstring")
                st.code(analysis_dict.get("docstring", "No docstring was generated."), language="python")
    else:
        st.warning("Please paste some code into the text area above.")