import gradio as gr
from langgraph.graph import StateGraph
from typing import TypedDict, Annotated, List, Dict
from langgraph.graph.message import add_messages
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
import json
import requests
import os
from dotenv import load_dotenv
import time

# Load environment variables
load_dotenv()

# Define the state structure
class State(TypedDict):
    messages: Annotated[list[SystemMessage | HumanMessage | AIMessage], add_messages]
    current_step: str
    code: str
    style_analysis: Dict
    security_analysis: Dict
    performance_analysis: Dict
    architecture_analysis: Dict
    final_recommendations: Dict

def call_huggingface_api(prompt: str, max_retries=3) -> Dict:
    """Call Hugging Face API with retry logic and proper error handling."""
    api_key = os.getenv("HUGGINGFACE_API_KEY")
    if not api_key:
        raise ValueError("HUGGINGFACE_API_KEY not found in environment variables")
    
    # You can change this to any model you prefer
    API_URL = "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1"
    headers = {"Authorization": f"Bearer {api_key}"}
    
    for attempt in range(max_retries):
        try:
            response = requests.post(
                API_URL,
                headers=headers,
                json={
                    "inputs": prompt,
                    "parameters": {
                        "max_new_tokens": 1000,
                        "temperature": 0.7,
                        "top_p": 0.95,
                        "return_full_text": False
                    }
                }
            )
            
            if response.status_code == 200:
                result = response.json()
                if isinstance(result, list) and len(result) > 0:
                    # Extract the generated text
                    text = result[0].get('generated_text', '')
                    # Try to parse as JSON if it contains JSON
                    try:
                        # Find JSON content between triple backticks if present
                        if "```json" in text:
                            json_str = text.split("```json")[1].split("```")[0].strip()
                        else:
                            json_str = text.strip()
                        return json.loads(json_str)
                    except json.JSONDecodeError:
                        return {"error": "Failed to parse JSON from response", "raw_text": text}
                
            # If model is loading, wait and retry
            if response.status_code == 503:
                wait_time = 2 ** attempt
                time.sleep(wait_time)
                continue
                
        except Exception as e:
            if attempt == max_retries - 1:
                return {"error": f"API call failed: {str(e)}"}
            time.sleep(2 ** attempt)
    
    return {"error": "Maximum retries reached"}

def analyze_code_style(state: State) -> dict:
    """Analyze code style and best practices."""
    code = state["code"]
    prompt = f"""You are a senior code reviewer focused on code style and best practices. Analyze this code:

{code}

Focus on:
1. Code readability and clarity
2. Adherence to common style guides
3. Variable/function naming
4. Code organization
5. Documentation quality

Provide your response in JSON format with these exact keys:
{{
    "issues": ["list of identified style issues"],
    "suggestions": ["list of improvement suggestions"],
    "overall_rating": "1-10 score as a number",
    "primary_concerns": ["list of main style concerns"]
}}"""

    analysis = call_huggingface_api(prompt)
    if "error" in analysis:
        analysis = {
            "issues": ["Error analyzing code style"],
            "suggestions": ["Try again later"],
            "overall_rating": 0,
            "primary_concerns": ["Analysis failed"]
        }
    
    messages = state["messages"] + [AIMessage(content="Completed code style analysis")]
    return {**state, "messages": messages, "style_analysis": analysis, "current_step": "security"}

def analyze_security(state: State) -> dict:
    """Analyze security vulnerabilities."""
    code = state["code"]
    prompt = f"""You are a security expert. Analyze this code for security vulnerabilities:

{code}

Focus on:
1. Input validation
2. Authentication/Authorization
3. Data exposure
4. Common vulnerabilities
5. Security best practices

Provide your response in JSON format with these exact keys:
{{
    "vulnerabilities": ["list of potential security issues"],
    "risk_levels": {{"vulnerability": "risk level"}},
    "recommendations": ["list of security improvements"],
    "overall_security_score": "1-10 score as a number"
}}"""

    analysis = call_huggingface_api(prompt)
    if "error" in analysis:
        analysis = {
            "vulnerabilities": ["Error analyzing security"],
            "risk_levels": {"Error": "High"},
            "recommendations": ["Try again later"],
            "overall_security_score": 0
        }
    
    messages = state["messages"] + [AIMessage(content="Completed security analysis")]
    return {**state, "messages": messages, "security_analysis": analysis, "current_step": "performance"}

def analyze_performance(state: State) -> dict:
    """Analyze code performance."""
    code = state["code"]
    prompt = f"""You are a performance optimization expert. Analyze this code for performance issues:

{code}

Focus on:
1. Time complexity
2. Space complexity
3. Resource usage
4. Bottlenecks
5. Optimization opportunities

Provide your response in JSON format with these exact keys:
{{
    "bottlenecks": ["list of identified performance bottlenecks"],
    "complexity_analysis": {{
        "time_complexity": "Big O notation",
        "space_complexity": "Big O notation",
        "critical_sections": ["list of critical sections"]
    }},
    "optimization_suggestions": ["list of performance improvements"],
    "performance_score": "1-10 score as a number"
}}"""

    analysis = call_huggingface_api(prompt)
    if "error" in analysis:
        analysis = {
            "bottlenecks": ["Error analyzing performance"],
            "complexity_analysis": {
                "time_complexity": "Unknown",
                "space_complexity": "Unknown",
                "critical_sections": []
            },
            "optimization_suggestions": ["Try again later"],
            "performance_score": 0
        }
    
    messages = state["messages"] + [AIMessage(content="Completed performance analysis")]
    return {**state, "messages": messages, "performance_analysis": analysis, "current_step": "architecture"}

def analyze_architecture(state: State) -> dict:
    """Analyze code architecture patterns."""
    code = state["code"]
    prompt = f"""You are a software architect. Analyze this code's architectural patterns:

{code}

Focus on:
1. Design patterns used
2. Code modularity
3. Component relationships
4. Architectural anti-patterns
5. System design principles

Provide your response in JSON format with these exact keys:
{{
    "patterns_identified": ["list of design patterns found"],
    "architectural_issues": ["list of architectural concerns"],
    "improvement_suggestions": ["list of architectural improvements"],
    "architecture_score": "1-10 score as a number"
}}"""

    analysis = call_huggingface_api(prompt)
    if "error" in analysis:
        analysis = {
            "patterns_identified": ["Error analyzing architecture"],
            "architectural_issues": ["Analysis failed"],
            "improvement_suggestions": ["Try again later"],
            "architecture_score": 0
        }
    
    messages = state["messages"] + [AIMessage(content="Completed architecture analysis")]
    return {**state, "messages": messages, "architecture_analysis": analysis, "current_step": "recommendations"}

def generate_final_recommendations(state: State) -> dict:
    """Generate final recommendations based on all analyses."""
    code = state["code"]
    prompt = f"""Analyze all previous results and provide final recommendations for this code:

Style Analysis: {json.dumps(state.get('style_analysis', {}))}
Security Analysis: {json.dumps(state.get('security_analysis', {}))}
Performance Analysis: {json.dumps(state.get('performance_analysis', {}))}
Architecture Analysis: {json.dumps(state.get('architecture_analysis', {}))}

Provide your response in JSON format with these exact keys:
{{
    "critical_issues": ["list of most critical issues"],
    "priority_improvements": ["list of high-priority improvements"],
    "quick_wins": ["list of easy-to-implement improvements"],
    "long_term_suggestions": ["list of long-term improvements"],
    "overall_health_score": "1-10 score as a number"
}}"""

    recommendations = call_huggingface_api(prompt)
    if "error" in recommendations:
        recommendations = {
            "critical_issues": ["Error generating recommendations"],
            "priority_improvements": ["Try again later"],
            "quick_wins": [],
            "long_term_suggestions": [],
            "overall_health_score": 0
        }
    
    messages = state["messages"] + [AIMessage(content="Generated final recommendations")]
    return {**state, "messages": messages, "final_recommendations": recommendations, "current_step": "end"}

def format_output(state: State) -> str:
    """Format the analysis results into a readable output."""
    output = """🔍 Code Analysis Report

🎨 Style & Best Practices
"""
    style = state.get("style_analysis", {})
    output += f"Rating: {style.get('overall_rating', 'N/A')}/10\n"
    output += "Issues:\n" + "\n".join([f"• {issue}" for issue in style.get("issues", [])]) + "\n\n"

    output += """🔒 Security Analysis
"""
    security = state.get("security_analysis", {})
    output += f"Score: {security.get('overall_security_score', 'N/A')}/10\n"
    vulnerabilities = security.get("vulnerabilities", [])
    risk_levels = security.get("risk_levels", {})
    output += "Vulnerabilities:\n" + "\n".join([f"• {v} ({risk_levels.get(v, 'Unknown')})" for v in vulnerabilities]) + "\n\n"

    output += """⚡ Performance Analysis
"""
    perf = state.get("performance_analysis", {})
    output += f"Score: {perf.get('performance_score', 'N/A')}/10\n"
    output += "Bottlenecks:\n" + "\n".join([f"• {b}" for b in perf.get("bottlenecks", [])]) + "\n\n"

    output += """🏗️ Architecture Analysis
"""
    arch = state.get("architecture_analysis", {})
    output += f"Score: {arch.get('architecture_score', 'N/A')}/10\n"
    output += "Patterns:\n" + "\n".join([f"• {p}" for p in arch.get("patterns_identified", [])]) + "\n\n"

    output += """📋 Final Recommendations
"""
    final = state.get("final_recommendations", {})
    output += f"Overall Health Score: {final.get('overall_health_score', 'N/A')}/10\n\n"
    output += "Critical Issues:\n" + "\n".join([f"• {i}" for i in final.get("critical_issues", [])]) + "\n\n"
    output += "Priority Improvements:\n" + "\n".join([f"• {i}" for i in final.get("priority_improvements", [])])

    return output

# Create and setup graph
workflow = StateGraph(State)

# Add nodes
workflow.add_node("style", analyze_code_style)
workflow.add_node("security", analyze_security)
workflow.add_node("performance", analyze_performance)
workflow.add_node("architecture", analyze_architecture)
workflow.add_node("recommendations", generate_final_recommendations)

# Add edges
workflow.add_edge("style", "security")
workflow.add_edge("security", "performance")
workflow.add_edge("performance", "architecture")
workflow.add_edge("architecture", "recommendations")

# Set entry and finish points
workflow.set_entry_point("style")
workflow.set_finish_point("recommendations")

# Compile the workflow
agent = workflow.compile()

def analyze_code(code: str) -> str:
    """Analyze the provided code using multiple perspectives."""
    initial_state = State(
        messages=[SystemMessage(content="Starting code analysis...")],
        current_step="style",
        code=code,
        style_analysis={},
        security_analysis={},
        performance_analysis={},
        architecture_analysis={},
        final_recommendations={}
    )
    
    final_state = agent.invoke(initial_state)
    return format_output(final_state)

# Create Gradio interface
iface = gr.Interface(
    fn=analyze_code,
    inputs=gr.Code(
        label="Enter your code for analysis",
        language="python",
        lines=20
    ),
    outputs=gr.Textbox(
        label="Analysis Results",
        lines=25
    ),
    title="🔍 Code Architecture Critic",
    description="Paste your code to get a comprehensive analysis of style, security, performance, and architecture.",
    examples=[
        ['''def process_data(data):
    result = []
    for i in range(len(data)):
        for j in range(len(data)):
            if data[i] + data[j] == 10:
                result.append((data[i], data[j]))
    return result

def save_to_db(user_input):
    query = "INSERT INTO users VALUES ('" + user_input + "')"
    db.execute(query)

API_KEY = "sk_test_123456789"''']
    ],
    theme=gr.themes.Soft()
)

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