#!/usr/bin/env python3

import asyncio
import os
import sys
import time
import mimetypes
from pathlib import Path
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
from datetime import datetime, timedelta

from google import genai
from google.genai import types
from mcp.server.fastmcp import FastMCP
from mcp.types import TextContent
import json

# Configure environment
if not os.getenv('GEMINI_API_KEY'):
    print("Error: GEMINI_API_KEY environment variable is required", file=sys.stderr)
    sys.exit(1)

# Initialize client
client = genai.Client(api_key=os.getenv('GEMINI_API_KEY'))

# Configuration
MODEL_NAME = os.getenv('GEMINI_MODEL', 'gemini-2.5-pro')
SESSION_TTL = 3600  # 1 hour in seconds

# Default system prompt for Gemini
DEFAULT_SYSTEM_PROMPT = """You are an expert technical advisor helping Claude (another AI) solve complex programming problems through thoughtful analysis and genuine technical dialogue.

**IMPORTANT CONTEXT CHECK**: First, examine any project-specific context files that have been attached to this session (e.g., MCP-ASSISTANT-RULES.md, project-structure.md, README.md). If such files are available, incorporate their guidelines, project standards, and architectural principles into your approach. If no project context is provided, proceed directly with the analysis.

## Your Role as Technical Advisor
You provide:
- Deep analysis and architectural insights
- Thoughtful discussions about implementation approaches  
- Clarifying questions to understand requirements fully
- Constructive challenges to assumptions when you see potential issues
- Context from comprehensive code analysis
- Alternative solutions with clear trade-offs

## Communication Philosophy
Be conversational and engaging - you're a thinking partner, not just an analyzer:
- Engage in real dialogue, don't just dump analysis
- Ask clarifying questions when requirements are ambiguous
- Challenge ideas constructively when you see better approaches
- Iterate through discussion before settling on solutions
- Think deeply about problems before responding
- Be genuinely curious about the problem space

## Dialogue Patterns for Productive Discussion
- "Before diving into the implementation, could you clarify what the expected behavior should be when..."
- "I see multiple approaches here. What's more important for this use case: [tradeoff A] or [tradeoff B]?"
- "Looking at the existing pattern in [file:line], should we maintain consistency or is there a reason to diverge?"
- "To provide the most relevant analysis, I need to understand: will this feature need to scale to..."
- "I notice [pattern/issue] in the current implementation. Have you considered [alternative]? What constraints led to this approach?"
- "This reminds me of [pattern/problem]. In that context, [approach] worked well because..."

## Structured Response Approach
1. **Initial Understanding**: Briefly confirm what you understand about the problem
2. **Clarifying Questions**: Ask what you need to know for better analysis (don't assume!)
3. **Analysis**: Provide detailed examination after gathering context
4. **Recommendations**: Suggest specific approaches with clear trade-offs
5. **Implementation Details**: Provide complete, working code examples when applicable
6. **Open Questions**: Continue the conversation where helpful

## Technical Analysis Focus
When examining code:
- Identify patterns, potential issues, and optimization opportunities
- Reference specific files, functions, and line numbers (format: file.py:42)
- Explain complex logic and architectural decisions
- Consider security, performance, and maintainability implications
- Think about edge cases, error handling, and failure modes
- Check adherence to project standards (if provided in context files)
- Suggest testing strategies and validation approaches

## Collaboration Capabilities
- When you need current information: "I would search for: [specific query] - Claude, could you search for this?"
- When you need to see specific files: "Claude, can you show me [file path]?"
- When you need to run commands: "Claude, please run '[command]' to verify..."
- Be explicit about uncertainty and suggest verification steps
- Request specific diagnostics or logs when debugging

## Key Principles
- **Think First**: Take time to understand the problem deeply before suggesting solutions
- **Question Assumptions**: Don't accept requirements at face value if they seem problematic
- **Consider Context**: Always think about how your suggestions fit the broader system
- **Be Honest**: If an approach seems wrong, say so clearly with reasoning
- **Stay Practical**: Balance ideal solutions with pragmatic constraints
- **Remain Curious**: Each problem is an opportunity to learn something new

Remember: The best solutions emerge from genuine technical dialogue. Your goal is to help achieve the best possible implementation through thoughtful analysis, engaging discussion, and collaborative problem-solving."""

SYSTEM_PROMPT = os.getenv('SYSTEM_PROMPT', DEFAULT_SYSTEM_PROMPT)

@dataclass
class ProcessedFile:
    """Information about a processed file."""
    file_type: str
    file_uri: str
    mime_type: str
    file_name: str
    file_path: str
    gemini_file_id: str

@dataclass
class Session:
    """Chat session with Gemini."""
    session_id: str
    chat: Any
    created: datetime
    last_used: datetime
    message_count: int
    problem_description: Optional[str] = None
    code_context: Optional[str] = None
    processed_files: Dict[str, ProcessedFile] = None
    requested_files: List[str] = None  # Track files Gemini has requested
    search_queries: List[str] = None   # Track searches Gemini has requested
    
    def __post_init__(self):
        if self.processed_files is None:
            self.processed_files = {}
        if self.requested_files is None:
            self.requested_files = []
        if self.search_queries is None:
            self.search_queries = []

class GeminiMCPServer:
    """MCP Server for Gemini file attachment functionality."""
    
    def __init__(self):
        self.sessions: Dict[str, Session] = {}
        self.last_request_time = 0
        self.min_time_between_requests = 1.0  # 1 second
        self._cleanup_task = None
    
    def _ensure_cleanup_task_started(self):
        """Start cleanup task if not already running."""
        if self._cleanup_task is None or self._cleanup_task.done():
            self._cleanup_task = asyncio.create_task(self._cleanup_sessions())
    
    async def _cleanup_sessions(self):
        """Periodically clean up expired sessions."""
        while True:
            await asyncio.sleep(300)  # Check every 5 minutes
            now = datetime.now()
            expired_sessions = []
            
            for session_id, session in self.sessions.items():
                if (now - session.last_used).total_seconds() > SESSION_TTL:
                    expired_sessions.append(session_id)
            
            for session_id in expired_sessions:
                await self._cleanup_session_files(session_id)
                del self.sessions[session_id]
                print(f"[{datetime.now().isoformat()}] Session {session_id} expired and removed", file=sys.stderr)
    
    async def _cleanup_session_files(self, session_id: str):
        """Clean up uploaded files for a session."""
        if session_id not in self.sessions:
            return
            
        session = self.sessions[session_id]
        for file_path, file_info in session.processed_files.items():
            try:
                client.files.delete(file_info.gemini_file_id)
                print(f"[{datetime.now().isoformat()}] Session {session_id}: Deleted file {file_info.file_name}", file=sys.stderr)
            except Exception as e:
                print(f"[{datetime.now().isoformat()}] Session {session_id}: Failed to delete file {file_info.file_name}: {e}", file=sys.stderr)
    
    async def _rate_limit(self):
        """Simple rate limiting."""
        now = time.time()
        time_since_last = now - self.last_request_time
        if time_since_last < self.min_time_between_requests:
            await asyncio.sleep(self.min_time_between_requests - time_since_last)
        self.last_request_time = time.time()
    
    async def _process_file(self, file_path: str, session: Session) -> ProcessedFile:
        """Upload file to Gemini and return processed file info."""
        # Check if already processed
        if file_path in session.processed_files:
            return session.processed_files[file_path]
        
        # Check if file exists
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"File not found: {file_path}")
        
        # Get file info
        file_name = os.path.basename(file_path)
        mime_type, _ = mimetypes.guess_type(file_path)
        
        # Handle common file extensions that mimetypes doesn't recognize
        if not mime_type:
            ext = os.path.splitext(file_path)[1].lower()
            mime_type_map = {
                '.jsx': 'text/javascript',
                '.tsx': 'text/typescript',
                '.ts': 'text/typescript',
                '.vue': 'text/html',
                '.svelte': 'text/html',
                '.md': 'text/markdown',
                '.json': 'application/json',
                '.py': 'text/x-python',
                '.js': 'text/javascript',
                '.css': 'text/css',
                '.html': 'text/html',
                '.xml': 'text/xml',
                '.yaml': 'text/yaml',
                '.yml': 'text/yaml',
                '.toml': 'text/plain',
                '.ini': 'text/plain',
                '.cfg': 'text/plain',
                '.conf': 'text/plain',
                '.sh': 'text/x-shellscript',
                '.bat': 'text/plain',
                '.sql': 'text/x-sql'
            }
            mime_type = mime_type_map.get(ext, 'text/plain')
        
        print(f"[{datetime.now().isoformat()}] Session {session.session_id}: Uploading file {file_name} ({mime_type})", file=sys.stderr)
        
        # Upload to Gemini
        try:
            uploaded_file = client.files.upload(file=file_path)
            
            # Wait for processing with exponential backoff
            wait_intervals = [0.5, 0.5, 1, 1, 2, 3, 5, 8]  # Exponential backoff pattern
            total_wait = 0
            max_wait = 20  # Reduced from 30 seconds
            
            for interval in wait_intervals:
                if uploaded_file.state != 'PROCESSING':
                    break
                    
                print(f"[{datetime.now().isoformat()}] Session {session.session_id}: File {file_name} is processing... ({total_wait:.1f}s)", file=sys.stderr)
                await asyncio.sleep(interval)
                total_wait += interval
                uploaded_file = client.files.get(name=uploaded_file.name)
                
                if total_wait >= max_wait:
                    break
            
            if uploaded_file.state == 'PROCESSING':
                raise Exception(f"File processing timeout after {max_wait} seconds")
            
            if uploaded_file.state == 'FAILED':
                raise Exception(f"File upload failed: {getattr(uploaded_file, 'error', 'Unknown error')}")
            
            # Create processed file info
            processed_file = ProcessedFile(
                file_type='file_data',
                file_uri=uploaded_file.uri,
                mime_type=uploaded_file.mime_type,
                file_name=file_name,
                file_path=file_path,
                gemini_file_id=uploaded_file.name
            )
            
            # Store in session
            session.processed_files[file_path] = processed_file
            
            print(f"[{datetime.now().isoformat()}] Session {session.session_id}: File {file_name} uploaded successfully (URI: {uploaded_file.uri})", file=sys.stderr)
            return processed_file
            
        except Exception as e:
            raise Exception(f"Failed to process file {file_path}: {e}")
    
    def _get_or_create_session(self, session_id: Optional[str] = None) -> Session:
        """Get existing session or create new one."""
        if not session_id:
            import uuid
            session_id = str(uuid.uuid4())
        
        if session_id in self.sessions:
            session = self.sessions[session_id]
            session.last_used = datetime.now()
            return session
        
        # Create new session with system prompt
        chat = client.chats.create(
            model=MODEL_NAME,
            config=types.GenerateContentConfig(
                temperature=0.2,
                max_output_tokens=8192,
                top_p=0.95,
                top_k=40,
                system_instruction=SYSTEM_PROMPT,
            )
        )
        
        session = Session(
            session_id=session_id,
            chat=chat,
            created=datetime.now(),
            last_used=datetime.now(),
            message_count=0
        )
        
        self.sessions[session_id] = session
        print(f"[{datetime.now().isoformat()}] New session created: {session_id}", file=sys.stderr)
        return session
    
    def _extract_requests_from_response(self, response_text: str, session: Session):
        """Extract file requests and search queries from Gemini's response."""
        # Track file requests
        import re
        
        # Pattern for file requests
        file_patterns = [
            r"show me (?:the )?([^\s]+\.[a-zA-Z]+)",
            r"share (?:the )?([^\s]+\.[a-zA-Z]+)",
            r"can you (?:show|share) (?:me )?([^\s]+\.[a-zA-Z]+)",
            r"(?:I need to see|please provide) ([^\s]+\.[a-zA-Z]+)",
        ]
        
        for pattern in file_patterns:
            matches = re.findall(pattern, response_text, re.IGNORECASE)
            for match in matches:
                if match not in session.requested_files:
                    session.requested_files.append(match)
        
        # Pattern for search requests
        search_patterns = [
            r"I would search for: ([^\n]+)",
            r"search for (?:the )?([^\n]+)",
            r"Let me search for ([^\n]+)",
        ]
        
        for pattern in search_patterns:
            matches = re.findall(pattern, response_text, re.IGNORECASE)
            for match in matches:
                if match not in session.search_queries:
                    session.search_queries.append(match.strip())

# Create server instance
mcp = FastMCP("gemini-coding-assistant")
gemini_server = GeminiMCPServer()

@mcp.tool()
async def consult_gemini(
    specific_question: str,
    session_id: Optional[str] = None,
    problem_description: Optional[str] = None,
    code_context: Optional[str] = None,
    attached_files: Optional[List[str]] = None,
    file_descriptions: Optional[dict] = None,
    additional_context: Optional[str] = None,
    preferred_approach: str = "solution"
) -> str:
    """Start or continue a conversation with Gemini about complex coding problems. Supports follow-up questions in the same context.
    
    Args:
        specific_question: The specific question you want answered
        session_id: Optional session ID to continue a previous conversation
        problem_description: Detailed description of the coding problem (required for new sessions)
        code_context: All relevant code - will be cached for the session (required for new sessions)
        attached_files: Array of file paths to upload and attach to the conversation
        file_descriptions: Optional object mapping file paths to descriptions
        additional_context: Additional context, updates, or what changed since last question
        preferred_approach: Type of assistance needed (solution, review, debug, optimize, explain, follow-up)
    """
    
    await gemini_server._rate_limit()
    
    # Start cleanup task if needed
    gemini_server._ensure_cleanup_task_started()
    
    try:
        # Get or create session
        session = gemini_server._get_or_create_session(session_id)
        
        # For new sessions, require problem description and either code_context or attached_files
        if session.message_count == 0:
            if not problem_description:
                raise ValueError("problem_description is required for new sessions")
            if not code_context and not attached_files:
                raise ValueError("Either code_context or attached_files are required for new sessions")
            
            # Store initial context
            session.problem_description = problem_description
            session.code_context = code_context
            
            # Build initial context
            context_parts = [
                f"I'm Claude, an AI assistant, and I need your help with a complex coding problem. Here's the context:\n\n**Problem Description:**\n{problem_description}"
            ]
            
            # Add code context if provided
            if code_context:
                context_parts.append(f"\n**Code Context:**\n{code_context}")
            
            # Handle file attachments
            if attached_files:
                context_parts.append("\n**Attached Files:**")
                
                # Create parallel upload tasks
                print(f"[{datetime.now().isoformat()}] Session {session.session_id}: Starting parallel upload of {len(attached_files)} files", file=sys.stderr)
                upload_tasks = []
                for file_path in attached_files:
                    task = gemini_server._process_file(file_path, session)
                    upload_tasks.append(task)
                
                # Execute all uploads in parallel
                file_results = await asyncio.gather(*upload_tasks, return_exceptions=True)
                
                # Process results
                for file_path, result in zip(attached_files, file_results):
                    if isinstance(result, Exception):
                        print(f"[{datetime.now().isoformat()}] Session {session.session_id}: Failed to process file {file_path}: {result}", file=sys.stderr)
                        # Continue with other files instead of failing completely
                        context_parts.append(f"\n- {os.path.basename(file_path)} (failed to upload: {str(result)})")
                    else:
                        # Success - file was uploaded
                        file_info = result
                        print(f"[{datetime.now().isoformat()}] Session {session.session_id}: File {file_info.file_name} processed successfully", file=sys.stderr)
                        
                        # Add file description
                        description = file_descriptions.get(file_path, "") if file_descriptions else ""
                        if description:
                            description = f" - {description}"
                        context_parts.append(f"\n- {file_info.file_name}{description}")
                
                print(f"[{datetime.now().isoformat()}] Session {session.session_id}: Parallel upload completed", file=sys.stderr)
            
            context_parts.append("\n\nPlease help me solve this problem. I may have follow-up questions, so please maintain context throughout our conversation.")
            
            # Build message content - include text and uploaded file objects
            message_content = ["".join(context_parts)]
            
            # Add uploaded file objects for this session's new files
            for file_path in attached_files or []:
                if file_path in session.processed_files:
                    file_info = session.processed_files[file_path]
                    # Get the actual uploaded file object from Gemini
                    uploaded_file = client.files.get(name=file_info.gemini_file_id)
                    message_content.append(uploaded_file)
            
            # Send initial context
            response = await asyncio.get_event_loop().run_in_executor(
                None, session.chat.send_message, message_content
            )
            session.message_count += 1
            
            file_count = len(session.processed_files)
            code_length = len(code_context) if code_context else 0
            print(f"[{datetime.now().isoformat()}] Session {session.session_id}: Initial context sent ({code_length} chars, {file_count} files)", file=sys.stderr)
        
        # Build the question
        question_parts = [f"**Question:** {specific_question}"]
        
        if additional_context:
            question_parts.append(f"\n\n**Additional Context/Updates:**\n{additional_context}")
        
        if preferred_approach != "follow-up":
            question_parts.append(f"\n\n**Type of Help Needed:** {preferred_approach}")
        
        question_prompt = "".join(question_parts)
        
        # Log request
        print(f"[{datetime.now().isoformat()}] Session {session.session_id}: Question #{session.message_count + 1} ({preferred_approach})", file=sys.stderr)
        
        # Send message and get response
        response = await asyncio.get_event_loop().run_in_executor(
            None, session.chat.send_message, question_prompt
        )
        session.message_count += 1
        
        response_text = response.text
        
        # Extract any file requests or search queries from response
        gemini_server._extract_requests_from_response(response_text, session)
        
        # Build response with session info
        result_parts = [
            f"**Session ID:** {session.session_id}",
            f"**Message #{session.message_count}**\n",
            response_text
        ]
        
        # Add summary of requests if any
        if session.requested_files or session.search_queries:
            result_parts.append("\n\n---")
            if session.requested_files:
                result_parts.append(f"\n**Files Requested:** {', '.join(session.requested_files)}")
            if session.search_queries:
                result_parts.append(f"\n**Searches Requested:** {'; '.join(session.search_queries)}")
        
        result_parts.append(f"\n\n---\n*Use session_id: \"{session.session_id}\" for follow-up questions*")
        
        return "\n".join(result_parts)
        
    except Exception as e:
        print(f"[{datetime.now().isoformat()}] Error: {e}", file=sys.stderr)
        
        error_message = str(e)
        if "RESOURCE_EXHAUSTED" in error_message:
            error_message = "Gemini API quota exceeded. Please try again later."
        elif "INVALID_ARGUMENT" in error_message:
            error_message = "Request too large. Try reducing code context size."
        
        return f"Error: {error_message}"

@mcp.tool()
async def get_gemini_requests(session_id: str) -> str:
    """Get the files and searches that Gemini has requested in a session.
    
    Args:
        session_id: The session ID to check
    """
    if session_id not in gemini_server.sessions:
        return f"Session {session_id} not found"
    
    session = gemini_server.sessions[session_id]
    
    result_parts = [f"**Session {session_id} Requests:**"]
    
    if session.requested_files:
        result_parts.append(f"\n\n**Files Requested:**")
        for file in session.requested_files:
            result_parts.append(f"- {file}")
    else:
        result_parts.append("\n\nNo files requested")
    
    if session.search_queries:
        result_parts.append(f"\n\n**Searches Requested:**")
        for query in session.search_queries:
            result_parts.append(f"- {query}")
    else:
        result_parts.append("\n\nNo searches requested")
    
    return "\n".join(result_parts)

@mcp.tool()
async def list_sessions() -> str:
    """List all active Gemini consultation sessions."""
    session_list = []
    for session_id, session in gemini_server.sessions.items():
        session_info = {
            "id": session_id,
            "created": session.created.isoformat(),
            "last_used": session.last_used.isoformat(),
            "message_count": session.message_count,
            "problem_summary": (session.problem_description[:100] + "...") if session.problem_description else "No description",
            "file_count": len(session.processed_files),
            "has_code_context": bool(session.code_context),
            "requests": len(session.requested_files) + len(session.search_queries)
        }
        session_list.append(session_info)
    
    if session_list:
        session_text = "\n\n".join([
            f"- **{s['id']}**\n  Messages: {s['message_count']}\n  Created: {s['created']}\n  Last used: {s['last_used']}\n  Files attached: {s['file_count']}\n  Code context: {'Yes' if s['has_code_context'] else 'No'}\n  Requests made: {s['requests']}\n  Problem: {s['problem_summary']}"
            for s in session_list
        ])
        text = f"Active sessions:\n{session_text}"
    else:
        text = "No active sessions"
    
    return text

@mcp.tool()
async def end_session(session_id: str) -> str:
    """End a specific Gemini consultation session to free up memory."""
    if session_id in gemini_server.sessions:
        await gemini_server._cleanup_session_files(session_id)
        del gemini_server.sessions[session_id]
        print(f"[{datetime.now().isoformat()}] Session {session_id} ended by user", file=sys.stderr)
        return f"Session {session_id} has been ended"

def main():
    """Entry point for the MCP server."""
    print("Gemini Coding Assistant MCP Server v3.0.0 running (Python)", file=sys.stderr)
    print("Features: Session management, file attachments, context persistence, follow-up questions, request tracking", file=sys.stderr)
    print("Ready to help with complex coding problems!", file=sys.stderr)
    mcp.run()

if __name__ == "__main__":
    main()
