import os
import json
import requests
from flask import Blueprint, request, jsonify
from flask_login import login_required

llm = Blueprint('llm', __name__)

def _call_llm(action, natural_query, code, runtime, logs):
    api_key = os.environ.get("GEMINI_API_KEY")
    model = os.environ.get("GEMINI_MODEL", "gemini-2.0-flash")
    
    if not api_key:
        return None, "GEMINI_API_KEY not set."

    api_url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
    headers = {'Content-Type': 'application/json'}

    system_instruction = ""
    user_prompt = ""

    if action == "generate":
        system_instruction = f"""
You are an expert {runtime} developer. Your task is to write a complete, production-ready {runtime} async function based on the user's request.

**Function Signature:**
Your function MUST have the following signature: `async (req, environment) => {{ ... }}` (for Node/Deno) or `def main(req, environment):` (for Python).
- `req`: An object containing details about the incoming HTTP request.
- `environment`: A mutable JSON object that persists across executions.

**Constraints:**
- The function will be executed in a sandboxed environment.
- Define ALL helper functions inside the main function.
- DO NOT use external module loading (require/import) unless standard lib. DO NOT IMPORT!
- Return ONLY the code.
"""
        user_prompt = f"Write a function that: {natural_query}"

    elif action == "explain":
        system_instruction = f"You are an expert {runtime} developer. Explain what the following code does in simple, concise terms suitable for a developer."
        user_prompt = f"Code:\n```\n{code}\n```"

    elif action == "optimize":
        system_instruction = f"""
You are an expert {runtime} developer. Analyze the following code and provide an optimized version.
Focus on performance, readability, and best practices.
Return ONLY the optimized code, without markdown formatting or explanation.
"""
        user_prompt = f"Code:\n```\n{code}\n```"

    elif action == "debug":
        system_instruction = f"""
You are an expert {runtime} developer. Analyze the following code and error logs.
Explain the error and provide a fixed version of the code.
Format your response as:
**Explanation:** [Brief explanation of the bug]
**Fix:**
```
[Fixed Code]
```
"""
        user_prompt = f"Code:\n```\n{code}\n```\n\nLogs:\n{logs}"

    try:
        payload = json.dumps({
            "contents": [{"parts": [{"text": system_instruction + "\n\n" + user_prompt}]}]
        })

        response = requests.post(api_url, headers=headers, data=payload)
        response.raise_for_status()

        response_data = response.json()
        
        candidates = response_data.get('candidates', [])
        if not candidates:
            return None, "No candidates found in API response."

        content = candidates[0].get('content', {})
        parts = content.get('parts', [])
        if not parts:
            return None, "No parts found in API response content."

        generated_text = parts[0].get('text', '').strip()

        # Clean up code blocks for generate/optimize actions
        if action in ["generate", "optimize"]:
            if generated_text.startswith("```"):
                # Find first newline
                first_newline = generated_text.find("\n")
                if first_newline != -1:
                    generated_text = generated_text[first_newline+1:]
            if generated_text.endswith("```"):
                generated_text = generated_text[:-3]
            generated_text = generated_text.strip()

        return generated_text, None

    except requests.exceptions.RequestException as e:
        return None, f"Error communicating with API: {e}"
    except (KeyError, IndexError, Exception) as e:
        return None, f"Error processing API response: {e}"

@llm.route("/generate_script", methods=["POST"])
@login_required
def generate_script():
    try:
        data = request.json
        action = data.get('action', 'generate') # Default to generate for backward compatibility
        natural_query = data.get('natural_query')
        code = data.get('code')
        runtime = data.get('runtime', 'node')
        logs = data.get('logs')

        if action == 'generate' and not natural_query:
             return jsonify({"error": "natural_query is required for generation"}), 400
        
        if action in ['explain', 'optimize', 'debug'] and not code:
            return jsonify({"error": "code is required for this action"}), 400

        result, error = _call_llm(action, natural_query, code, runtime, logs)

        if error:
            return jsonify({"error": error}), 500

        return jsonify({"result": result})

    except Exception as e:
        return jsonify({'error': str(e)}), 500