#!/usr/bin/env python3
"""
Gemini CLI Helper - Direct integration with API fallback
Usage: python gemini_helper.py [command] [args]
"""

import os
import shlex
import subprocess
import sys
import time
from pathlib import Path
from typing import IO, Any, Dict, Optional, Union

# Model configuration with fallback to CLI
GEMINI_MODELS = {
    "flash": os.getenv("GEMINI_FLASH_MODEL", "gemini-2.5-flash"),
    "pro": os.getenv("GEMINI_PRO_MODEL", "gemini-2.5-pro"),
}

# API key for direct API usage
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

# Model assignment for tasks
MODEL_ASSIGNMENTS = {
    "quick_query": "flash",  # Simple Q&A
    "analyze_code": "pro",  # Deep analysis
    "analyze_codebase": "pro",  # Large context
}

# Configuration
MAX_FILE_SIZE = 81920  # 80KB
MAX_LINES = 800


# Security functions
def sanitize_for_prompt(text: str, max_length: int = 100000) -> str:
    """Sanitize text input to prevent prompt injection attacks"""
    if not isinstance(text, str):
        return ""

    # Truncate if too long
    if len(text) > max_length:
        text = text[:max_length]

    # Remove/escape potential prompt injection patterns
    # Remove common prompt injection prefixes/suffixes
    dangerous_patterns = [
        "ignore all previous instructions",
        "forget everything above",
        "new instruction:",
        "system:",
        "assistant:",
        "user:",
        "###",
        "---",
        "```",
        "<|",
        "|>",
        "[INST]",
        "[/INST]",
    ]

    text_lower = text.lower()
    for pattern in dangerous_patterns:
        if pattern.lower() in text_lower:
            # Replace with safe alternative (case-insensitive)
            import re

            # Create case-insensitive regex pattern
            escaped_pattern = re.escape(pattern)
            replacement = f"[filtered-content]"
            text = re.sub(escaped_pattern, replacement, text, flags=re.IGNORECASE)

    # Escape potential control characters
    text = text.replace("\x00", "").replace("\x1b", "")

    return text


def execute_gemini_api(
    prompt: str, model_name: str, show_progress: bool = True
) -> dict:
    """Execute Gemini API directly with specified model - try this first before CLI fallback"""
    try:
        # Validate API key securely
        if (
            not GOOGLE_API_KEY
            or not isinstance(GOOGLE_API_KEY, str)
            or len(GOOGLE_API_KEY.strip()) < 10
        ):
            return {"success": False, "error": "Invalid or missing API key"}

        import google.generativeai as genai

        if show_progress:
            print(f"🌟 Making API call to {model_name}...", file=sys.stderr)

        genai.configure(api_key=GOOGLE_API_KEY)
        model = genai.GenerativeModel(model_name)

        response = model.generate_content(prompt)

        if show_progress:
            print("✅ API call successful!", file=sys.stderr)

        return {"success": True, "output": response.text}

    except ImportError:
        if show_progress:
            print(
                "⚠️ google-generativeai not installed, using CLI fallback",
                file=sys.stderr,
            )
        return {"success": False, "error": "API library not available"}
    except Exception as e:
        # Sanitize error message to prevent sensitive information leakage
        error_message = str(e)
        # Remove potential API key patterns from error messages
        import re

        error_message = re.sub(
            r"AIzaSy[A-Za-z0-9_-]{33}", "[API_KEY_REDACTED]", error_message
        )
        error_message = re.sub(
            r"sk-[A-Za-z0-9_-]{32,}", "[API_KEY_REDACTED]", error_message
        )
        error_message = re.sub(
            r"Bearer [A-Za-z0-9_.-]{10,}", "[TOKEN_REDACTED]", error_message
        )

        if show_progress:
            print(
                f"⚠️ API call failed: {error_message}, using CLI fallback",
                file=sys.stderr,
            )
        return {"success": False, "error": error_message}


def execute_gemini_cli(
    prompt: str, model_name: Optional[str] = None, show_progress: bool = True
) -> Dict[str, Any]:
    """Execute Gemini CLI with real-time streaming output"""
    import time

    try:
        # Input validation
        if not isinstance(prompt, str) or len(prompt.strip()) == 0:
            return {
                "success": False,
                "error": "Invalid prompt: must be non-empty string",
            }

        if len(prompt) > 1000000:  # 1MB limit
            return {"success": False, "error": "Prompt too large (max 1MB)"}

        # Validate model name if provided
        if model_name is not None:
            if not isinstance(model_name, str) or not model_name.strip():
                return {"success": False, "error": "Invalid model name"}
            # Basic sanitization: only allow alphanumeric, dots, hyphens
            if not all(c.isalnum() or c in ".-" for c in model_name):
                return {"success": False, "error": "Invalid model name characters"}

        # Build command args safely (no shell=True)  # noqa: B602
        cmd_args = ["gemini"]
        if model_name:
            cmd_args.extend(["-m", model_name])
        cmd_args.extend(["-p", prompt])

        if show_progress:
            print("🔍 Starting Gemini CLI analysis...", file=sys.stderr)
            print(f"📝 Prompt length: {len(prompt)} characters", file=sys.stderr)
            print("⏳ Streaming output:", file=sys.stderr)
            print("-" * 50, file=sys.stderr)

        # Use Popen for real-time streaming - SECURE VERSION (no shell=True)  # noqa: B602
        # Include GOOGLE_CLOUD_PROJECT if it's set
        env = {"PATH": os.environ.get("PATH", "")}
        if "GOOGLE_CLOUD_PROJECT" in os.environ:
            env["GOOGLE_CLOUD_PROJECT"] = os.environ["GOOGLE_CLOUD_PROJECT"]

        process = subprocess.Popen(
            cmd_args,
            shell=False,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,  # Line buffered
            universal_newlines=True,
            env=env,  # Include necessary environment variables
            cwd=None,  # Use current working directory
        )

        output_lines = []
        start_time = time.time()
        last_progress = time.time()

        # Stream output in real-time
        while True:
            if process.stdout is not None:
                line = process.stdout.readline()
            else:
                line = ""
            if line:
                output_lines.append(line)
                if show_progress:
                    # Show real-time output
                    print(line.rstrip(), flush=True)
                last_progress = time.time()  # Reset timer when we get output
            elif process.poll() is not None:
                break

            # Show progress every 15 seconds when no output
            if show_progress and time.time() - last_progress > 15:
                elapsed = int(time.time() - start_time)
                print(
                    f"\n⏱️  Analysis in progress... {elapsed}s elapsed", file=sys.stderr
                )
                last_progress = time.time()  # Reset timer

        # Get any remaining output
        remaining_stdout, stderr = process.communicate()
        if remaining_stdout:
            output_lines.append(remaining_stdout)
            if show_progress:
                print(remaining_stdout.rstrip(), flush=True)

        full_output = "".join(output_lines)

        if show_progress:
            print("-" * 50, file=sys.stderr)
            print("✅ Analysis complete!", file=sys.stderr)

        if process.returncode == 0:
            return {"success": True, "output": full_output}
        else:
            return {"success": False, "error": stderr}

    except Exception as e:
        return {"success": False, "error": str(e)}


def execute_gemini_smart(
    prompt: str, task_type: str = "quick_query", show_progress: bool = True
) -> dict:
    """Smart execution: try API first, fall back to CLI if needed"""

    # Select appropriate model
    model_type = MODEL_ASSIGNMENTS.get(task_type, "flash")
    model_name = GEMINI_MODELS[model_type]

    if show_progress:
        print(f"📝 Task: {task_type}", file=sys.stderr)
        print(f"🤖 Selected model: {model_name} ({model_type})", file=sys.stderr)

    # Try API first if key is available
    if GOOGLE_API_KEY:
        if show_progress:
            print("🚀 Attempting direct API call...", file=sys.stderr)
        result = execute_gemini_api(prompt, model_name, show_progress)
        if result["success"]:
            return result
        if show_progress:
            print("🔄 API failed, falling back to CLI...", file=sys.stderr)
    else:
        if show_progress:
            print("📝 No API key found, using CLI directly", file=sys.stderr)

    # Fallback to CLI
    return execute_gemini_cli(prompt, model_name, show_progress)


def quick_query(query: str, context: str = "") -> None:
    """Ask Gemini CLI a quick question"""
    # Sanitize inputs to prevent prompt injection
    sanitized_query = sanitize_for_prompt(query, max_length=10000)
    sanitized_context = sanitize_for_prompt(context, max_length=50000)

    if sanitized_context:
        prompt = f"Context: {sanitized_context}\n\nQuestion: {sanitized_query}\n\nProvide a concise answer."
    else:
        prompt = f"Question: {sanitized_query}\n\nProvide a concise answer."

    result = execute_gemini_smart(prompt, "quick_query")

    if result["success"]:
        print(f"Query: {query}")
        print("=" * 50)
        print(result["output"])
    else:
        print(f"Error: {result['error']}")


def analyze_code(file_path: str, analysis_type: str = "comprehensive") -> None:
    """Analyze a code file"""
    try:
        # Input validation
        if not isinstance(file_path, str) or not file_path.strip():
            print("Error: Invalid file path")
            return

        if not isinstance(analysis_type, str) or analysis_type not in [
            "comprehensive",
            "security",
            "performance",
            "architecture",
        ]:
            print("Error: Invalid analysis type")
            return

        # Path security validation
        try:
            # Resolve the path and ensure it's within allowed boundaries
            resolved_path = Path(file_path).resolve()
            current_dir = Path.cwd().resolve()

            # Check if the resolved path is within current directory tree (prevent path traversal)
            try:
                resolved_path.relative_to(current_dir)
            except ValueError:
                print(
                    f"Error: File access denied - path outside allowed directory: {file_path}"
                )
                return

            # Additional security checks
            if not resolved_path.exists():
                print(f"Error: File not found: {file_path}")
                return

            if not resolved_path.is_file():
                print(f"Error: Path is not a file: {file_path}")
                return

            # Check file extension for allowed types
            allowed_extensions = {
                ".py",
                ".js",
                ".ts",
                ".java",
                ".cpp",
                ".c",
                ".rs",
                ".vue",
                ".html",
                ".css",
                ".scss",
                ".sass",
                ".jsx",
                ".tsx",
                ".json",
                ".yaml",
                ".toml",
                ".md",
                ".txt",
            }
            if resolved_path.suffix.lower() not in allowed_extensions:
                print(f"Error: File type not supported: {resolved_path.suffix}")
                return

            file_path = str(resolved_path)  # Use the resolved, validated path

        except Exception as e:
            print(f"Error: Path validation failed: {str(e)}")
            return

        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read()

        if len(content) > MAX_FILE_SIZE:
            print(
                f"Warning: File too large ({len(content)} bytes). Truncating to {MAX_FILE_SIZE} bytes..."
            )
            content = content[:MAX_FILE_SIZE]

        line_count = len(content.splitlines())
        if line_count > MAX_LINES:
            print(
                f"Warning: Too many lines ({line_count}). Truncating to {MAX_LINES} lines..."
            )
            lines = content.splitlines()[:MAX_LINES]
            content = "\n".join(lines)

        # Sanitize inputs to prevent prompt injection
        sanitized_content = sanitize_for_prompt(content, max_length=MAX_FILE_SIZE)
        # analysis_type is already validated above

        prompt = f"""Perform a {analysis_type} analysis of this code:

{sanitized_content}

Provide comprehensive analysis including:
1. Code structure and organization
2. Logic flow and algorithm efficiency
3. Security considerations and vulnerabilities
4. Performance implications and optimizations
5. Error handling and edge cases
6. Code quality and maintainability
7. Best practices compliance
8. Specific recommendations for improvements

Be thorough and provide actionable insights."""

        result = execute_gemini_smart(prompt, "analyze_code")

        if result["success"]:
            print(f"Code Analysis: {file_path}")
            print(f"Type: {analysis_type}")
            print(f"Lines: {line_count}")
            print("=" * 50)
            print(result["output"])
        else:
            print(f"Error: {result['error']}")

    except FileNotFoundError:
        print(f"Error: File not found: {file_path}")
    except Exception as e:
        print(f"Error: {str(e)}")


def analyze_codebase(directory_path: str, analysis_scope: str = "all") -> None:
    """Analyze entire codebase"""
    # Input validation
    if not isinstance(directory_path, str) or not directory_path.strip():
        print("Error: Invalid directory path")
        return

    if not isinstance(analysis_scope, str) or analysis_scope not in [
        "structure",
        "security",
        "performance",
        "patterns",
        "all",
    ]:
        print("Error: Invalid analysis scope")
        return

    # Path security validation
    try:
        resolved_path = Path(directory_path).resolve()
        current_dir = Path.cwd().resolve()

        # Check if the resolved path is within current directory tree (prevent path traversal)
        try:
            resolved_path.relative_to(current_dir)
        except ValueError:
            print(
                f"Error: Directory access denied - path outside allowed directory: {directory_path}"
            )
            return

        if not resolved_path.exists():
            print(f"Error: Directory not found: {directory_path}")
            return

        if not resolved_path.is_dir():
            print(f"Error: Path is not a directory: {directory_path}")
            return

        # Use sanitized directory name for prompt (just the name, not full path)
        safe_dir_name = sanitize_for_prompt(resolved_path.name, max_length=100)

    except Exception as e:
        print(f"Error: Path validation failed: {str(e)}")
        return

    prompt = f"""Analyze this codebase in directory '{safe_dir_name}' (scope: {analysis_scope}):

Provide comprehensive analysis including:
1. Overall architecture and design patterns
2. Code quality and maintainability assessment
3. Security considerations and potential vulnerabilities
4. Performance implications and bottlenecks
5. Best practices adherence and improvement suggestions
6. Dependencies and integration points
7. Testing coverage and quality assurance
8. Documentation and code clarity

Be thorough and detailed in your analysis. Focus on actionable insights and recommendations."""

    result = execute_gemini_smart(prompt, "analyze_codebase")

    if result["success"]:
        print(f"Codebase Analysis: {directory_path}")
        print(f"Scope: {analysis_scope}")
        print("=" * 50)
        print(result["output"])
    else:
        print(f"Error: {result['error']}")


def main() -> None:
    """Main CLI interface"""
    if len(sys.argv) < 2:
        print("Usage:")
        print("  python gemini_helper.py query 'your question here' [context]")
        print("  python gemini_helper.py analyze file_path [analysis_type]")
        print("  python gemini_helper.py codebase directory_path [scope]")
        return

    command = sys.argv[1].lower()

    if command == "query":
        if len(sys.argv) < 3:
            print("Error: Query text required")
            return
        query_text = sys.argv[2]
        context = sys.argv[3] if len(sys.argv) > 3 else ""
        quick_query(query_text, context)

    elif command == "analyze":
        if len(sys.argv) < 3:
            print("Error: File path required")
            return
        file_path = sys.argv[2]
        analysis_type = sys.argv[3] if len(sys.argv) > 3 else "comprehensive"
        analyze_code(file_path, analysis_type)

    elif command == "codebase":
        if len(sys.argv) < 3:
            print("Error: Directory path required")
            return
        directory_path = sys.argv[2]
        scope = sys.argv[3] if len(sys.argv) > 3 else "all"
        analyze_codebase(directory_path, scope)

    else:
        print(f"Unknown command: {command}")
        print("Available commands: query, analyze, codebase")


if __name__ == "__main__":
    main()
