#!/usr/bin/env python3
"""
MCP Client for connecting to problems server via STDIO
Provides interface to call server methods via JSON-RPC over subprocess
"""
import json
import subprocess
import logging
import os
import sys
import io
from typing import Dict, List, Any, Optional

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    stream=sys.stderr  # Ensure logs go to stderr
)
logger = logging.getLogger(__name__)

# Make paths robust by defining them relative to this script's location
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
_SERVER_SCRIPT_PATH = os.path.join(_SCRIPT_DIR, "mcp_server.py")


class MCPClient:
    """MCP Client for problems server via STDIO"""
    
    def __init__(self, server_script: str = _SERVER_SCRIPT_PATH):
        self.server_script = server_script
        self.server_process: Optional[subprocess.Popen] = None
        self.request_id = 1
        
    def _start_server(self) -> bool:
        """Start the server process"""
        try:
            # Check if server script exists
            if not os.path.exists(self.server_script):
                logger.error(f"Server script not found: {self.server_script}")
                return False
            
            # Prepare environment for UTF-8 on Windows
            env = os.environ.copy()
            if sys.platform == "win32":
                env['PYTHONIOENCODING'] = 'utf-8'
            
            # Start server process with explicit UTF-8 handling
            self.server_process = subprocess.Popen(
                [sys.executable, self.server_script],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                text=True,
                encoding='utf-8',
                errors='replace',  # Replace problematic characters instead of failing
                bufsize=1,  # Line buffered
                env=env
            )
            
            logger.info(f"Started server process with PID: {self.server_process.pid}")
            
            # Wait for the server's readiness signal
            ready_line = self.server_process.stdout.readline()
            if "SERVER_READY" not in ready_line:
                logger.error(f"Server failed to send ready signal. Got: {ready_line.strip()}")
                self._stop_server()
                return False
            
            logger.info("Server is ready.")
            return True
            
        except Exception as e:
            logger.error(f"Failed to start server: {e}")
            return False
    
    def _stop_server(self):
        """Stop the server process"""
        if self.server_process:
            try:
                self.server_process.terminate()
                self.server_process.wait(timeout=5)
                logger.info("Server process stopped")
            except subprocess.TimeoutExpired:
                logger.warning("Server process did not stop gracefully, killing...")
                self.server_process.kill()
                self.server_process.wait()
            except Exception as e:
                logger.error(f"Error stopping server: {e}")
            finally:
                self.server_process = None
    
    def _make_request(self, method: str, params: Dict[str, Any] = None) -> Dict[str, Any]:
        """Make JSON-RPC request to server via STDIO"""
        try:
            # Ensure server is running
            if not self.server_process or self.server_process.poll() is not None:
                if not self._start_server():
                    return {
                        'success': False,
                        'error': 'Failed to start server'
                    }
            
            # Prepare request
            request_data = {
                'jsonrpc': '2.0',
                'method': method,
                'params': params or {},
                'id': self.request_id
            }
            self.request_id += 1
            
            logger.debug(f"Making request: {method} with params: {params}")
            
            # Send request
            request_line = json.dumps(request_data) + '\n'
            self.server_process.stdin.write(request_line)
            self.server_process.stdin.flush()
            
            # Read response
            response_line = self.server_process.stdout.readline()
            if not response_line:
                return {
                    'success': False,
                    'error': 'No response from server'
                }
            
            response_data = json.loads(response_line.strip())
            
            # Check for errors
            if 'error' in response_data:
                logger.error(f"Server error: {response_data['error']}")
                return {
                    'success': False,
                    'error': response_data['error']['message'],
                    'error_code': response_data['error']['code']
                }
            
            return {
                'success': True,
                'data': response_data['result']
            }
            
        except json.JSONDecodeError as e:
            logger.error(f"JSON decode error: {e}")
            return {
                'success': False,
                'error': f"Invalid JSON response: {str(e)}"
            }
        except Exception as e:
            logger.error(f"Request failed: {e}")
            return {
                'success': False,
                'error': f"Request error: {str(e)}"
            }
    
    def __enter__(self):
        """Context manager entry"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit"""
        self._stop_server()
    
    def health_check(self) -> Dict[str, Any]:
        """Check if server is healthy by making a simple request"""
        try:
            # Try to get first question as a health check (using physics grade 8 as default)
            result = self._make_request('get_question_from_db', {'subject_id': 1, 'grade': 8, 'question_order': 1})
            if result['success']:
                return {
                    'success': True,
                    'data': {'status': 'healthy'}
                }
            else:
                return result
        except Exception as e:
            logger.error(f"Health check failed: {e}")
            return {
                'success': False,
                'error': str(e)
            }
    

    


    def get_question_from_db(self, subject_id: int, grade: int, question_order: int) -> Dict[str, Any]:
        """Get specific question from database by subject_id, grade and question order"""
        logger.info(f"Requesting question {question_order} for subject {subject_id}, grade {grade}")
        result = self._make_request('get_question_from_db', {
            'subject_id': subject_id,
            'grade': grade,
            'question_order': question_order
        })
        
        if result['success']:
            data = result['data']
            subject_folder = data.get('subject_folder', 'unknown')
            logger.info(f"Got question {question_order} from {subject_folder}/{grade}")
            logger.info(f"Question: {data.get('question', '')[:100]}...")
        
        return result
    
    def list_available_questions(self, subject_id: int, grade: int) -> Dict[str, Any]:
        """List all available questions for a subject and grade from database"""
        logger.info(f"Requesting list of questions for subject {subject_id}, grade {grade}")
        result = self._make_request('list_available_questions', {
            'subject_id': subject_id,
            'grade': grade
        })
        
        if result['success']:
            data = result['data']
            subject_folder = data.get('subject_folder', 'unknown')
            total_questions = data.get('total_questions', 0)
            logger.info(f"Found {total_questions} questions for {subject_folder}/{grade}")
        
        return result

    def list_available_topics(self, subject_id: int) -> Dict[str, Any]:
        """List all available topics for a subject (grade-agnostic)"""
        logger.info(f"Requesting list of topics for subject {subject_id}")
        result = self._make_request('list_available_topics', {
            'subject_id': subject_id
        })
        
        if result['success']:
            data = result['data']
            subject_folder = data.get('subject_folder', 'unknown')
            total_topics = data.get('total_topics', 0)
            logger.info(f"Found {total_topics} topics for {subject_folder}")
        
        return result

    def save_paragraph_progress(self, user_id: int, textbook_id: int, seq_number: str, progress: str) -> Dict[str, Any]:
        """Save student progress for a specific paragraph to tu_paragraph_progress table"""
        logger.info(f"Saving progress for user {user_id}, textbook {textbook_id}, section {seq_number}")
        result = self._make_request('save_paragraph_progress', {
            'user_id': user_id,
            'textbook_id': textbook_id,
            'seq_number': seq_number,
            'progress': progress
        })
        
        if result['success']:
            data = result['data']
            section_title = data.get('section_title', 'unknown')
            logger.info(f"Progress saved successfully for section: {section_title}")
        
        return result

    def get_problems_by_topic(self, subject_id: int, topic_id: int, user_id: int) -> Dict[str, Any]:
        """Get problems by topic ID for a specific subject (no grade filtering)"""
        logger.info(f"Requesting problems for topic ID '{topic_id}' in subject {subject_id}")
        result = self._make_request('get_problems_by_topic', {
            'subject_id': subject_id,
            'topic_id': topic_id,
            'user_id': user_id
        })
        
        if result['success']:
            data = result['data']
            subject_folder = data.get('subject_folder', 'unknown')
            topic_name = data.get('topic_name', 'unknown')
            total_problems = data.get('total_problems', 0)
            logger.info(f"Found {total_problems} problems for topic ID: {topic_id}) in {subject_folder}")
        
        return result

    def mark_problem_as_solved(self, user_id: int, problem_id: int) -> Dict[str, Any]:
        """Mark a problem as solved for a specific user"""
        logger.info(f"Marking problem {problem_id} as solved for user {user_id}")
        result = self._make_request('mark_problem_as_solved', {
            'user_id': user_id,
            'problem_id': problem_id
        })
        
        if result['success']:
            logger.info(f"Problem {problem_id} marked as solved for user {user_id}")
            
        return result

# === LLM INTEGRATION FUNCTIONS ===

# Global client instance for LLM function calls
_global_client = None

def get_client():
    """Get or create global MCP client instance"""
    global _global_client
    if _global_client is None:
        _global_client = MCPClient()
    return _global_client





def llm_get_question_from_db(subject_id: int, grade: int, question_order: int) -> str:
    """
    LLM Function: Get specific question from database by subject_id, grade and question order
    
    Args:
        subject_id: ID of the subject
        grade: Grade as integer (7, 8, 9, 10, 11)
        question_order: Order number of the question to retrieve
        
    Returns:
        JSON string with question, answer, progress info, and next question info
    """
    try:
        # Convert subject_id to int if it's a string (from Gemini API)
        if isinstance(subject_id, str):
            try:
                subject_id = int(subject_id)
            except ValueError:
                return json.dumps({
                    'status': 'error',
                    'message': 'subject_id must be a valid integer'
                }, ensure_ascii=False, indent=2)
        
        if not isinstance(subject_id, int):
            return json.dumps({
                'status': 'error',
                'message': 'subject_id must be an integer'
            }, ensure_ascii=False, indent=2)
        
        # Convert grade to int if it's a string
        if isinstance(grade, str):
            try:
                grade = int(grade)
            except ValueError:
                return json.dumps({
                    'status': 'error',
                    'message': 'grade must be a valid integer'
                }, ensure_ascii=False, indent=2)
        
        if not isinstance(grade, int):
            return json.dumps({
                'status': 'error',
                'message': 'grade must be an integer'
            }, ensure_ascii=False, indent=2)
        
        # Convert question_order to int if it's a string
        if isinstance(question_order, str):
            try:
                question_order = int(question_order)
            except ValueError:
                return json.dumps({
                    'status': 'error',
                    'message': 'question_order must be a valid integer'
                }, ensure_ascii=False, indent=2)
        
        if not isinstance(question_order, int):
            return json.dumps({
                'status': 'error',
                'message': 'question_order must be an integer'
            }, ensure_ascii=False, indent=2)
        
        # Проверяем что grade в правильном диапазоне
        valid_grades = [7, 8, 9, 10, 11]
        if grade not in valid_grades:
            return json.dumps({
                'status': 'error',
                'message': f'Неверный класс: {grade}. Используй ТОЛЬКО: {", ".join(map(str, valid_grades))}'
            }, ensure_ascii=False, indent=2)
        
        client = get_client()
        result = client.get_question_from_db(subject_id, grade, question_order)
        
        if result['success']:
            data = result['data']
            
            # Prepare next question info
            next_question_order = question_order + 1 if not data['is_last_question'] else None
            
            response_data = {
                'status': 'success',
                'subject_id': data['subject_id'],
                'subject_folder': data['subject_folder'],
                'grade': data['grade'],
                'collection_name': data['collection_name'],
                'current_question': {
                    'question_order': data['question_order'],
                    'question': data['question'],
                    'answer': data['answer'],
                    'problem_id': data['problem_id']
                },
                'progress': {
                    'current_question_number': data['question_order'],
                    'total_questions': data['total_questions'],
                    'is_last_question': data['is_last_question'],
                    'questions_remaining': data['total_questions'] - data['question_order']
                },
                'next_action': {
                    'has_next_question': not data['is_last_question'],
                    'next_question_order': next_question_order,
                    'instruction': 'Запроси следующий вопрос с question_order=' + str(next_question_order) if next_question_order else 'Это был последний вопрос в тесте'
                }
            }
            
            return json.dumps(response_data, ensure_ascii=False, indent=2)
        else:
            return json.dumps({
                'status': 'error',
                'message': result['error']
            }, ensure_ascii=False, indent=2)
            
    except Exception as e:
        return json.dumps({
            'status': 'error',
            'message': f'Failed to get question: {str(e)}'
        }, ensure_ascii=False, indent=2)

def llm_list_available_questions(subject_id: int, grade: int) -> str:
    """
    [DEPRECATED] Use llm_list_available_topics instead.
    """
    try:
        return llm_list_available_topics(subject_id)
    except Exception as e:
        return json.dumps({'status': 'error', 'message': f'Deprecated: {str(e)}'}, ensure_ascii=False, indent=2)

def llm_list_available_topics(subject_id: int) -> str:
    """
    LLM Function: List all available topics for a subject (no grade filtering)
    
    Args:
        subject_id: ID of the subject
        
    Returns:
        JSON string with list of available topics or error message
    """
    try:
        # Convert subject_id to int if it's a string (from Gemini API)
        if isinstance(subject_id, str):
            try:
                subject_id = int(subject_id)
            except ValueError:
                return json.dumps({
                    'status': 'error',
                    'message': 'subject_id must be a valid integer'
                }, ensure_ascii=False, indent=2)
        
        if not isinstance(subject_id, int):
            return json.dumps({
                'status': 'error',
                'message': 'subject_id must be an integer'
            }, ensure_ascii=False, indent=2)
        
        client = get_client()
        result = client.list_available_topics(subject_id)
        
        if result['success']:
            data = result['data']
            return json.dumps({
                'status': 'success',
                'subject_id': data['subject_id'],
                'subject_folder': data['subject_folder'],
                'total_topics': data['total_topics'],
                'topics': data['topics']
            }, ensure_ascii=False, indent=2)
        else:
            return json.dumps({
                'status': 'error',
                'message': result['error']
            }, ensure_ascii=False, indent=2)
            
    except Exception as e:
        return json.dumps({
            'status': 'error',
            'message': f'Failed to list topics: {str(e)}'
        }, ensure_ascii=False, indent=2)

def llm_save_paragraph_progress(user_id: int, textbook_id: int, seq_number: str, progress: str) -> str:
    """
    LLM Function: Save student progress for a specific paragraph to tu_paragraph_progress table
    
    Args:
        user_id: ID of the user
        textbook_id: ID of the textbook
        seq_number: Section number from contents table (e.g., "1.1", "2.3", etc.)
        progress: Progress text describing what the student learned
        
    Returns:
        JSON string with success/error status and details
    """
    try:
        # Use constants for user_id and textbook_id
        # user_id = 1
        # textbook_id = 135
        
        # Validate user_id is an integer
        if not isinstance(user_id, int):
            return json.dumps({
                'status': 'error',
                'message': 'user_id must be an integer'
            }, ensure_ascii=False, indent=2)
            
        # Validate textbook_id is an integer
        if not isinstance(textbook_id, int):
            return json.dumps({
                'status': 'error',
                'message': 'textbook_id must be an integer'
            }, ensure_ascii=False, indent=2)
        
        # Validate seq_number is a string
        if not isinstance(seq_number, str) or not seq_number.strip():
            return json.dumps({
                'status': 'error',
                'message': 'seq_number must be a non-empty string'
            }, ensure_ascii=False, indent=2)
        
        # Validate progress is a string
        if not isinstance(progress, str) or not progress.strip():
            return json.dumps({
                'status': 'error',
                'message': 'progress must be a non-empty string'
            }, ensure_ascii=False, indent=2)
        
        # Limit progress text length
        if len(progress) > 1000:
            return json.dumps({
                'status': 'error',
                'message': 'progress text must be 1000 characters or less'
            }, ensure_ascii=False, indent=2)
        
        client = get_client()
        result = client.save_paragraph_progress(user_id, textbook_id, seq_number, progress)
        
        if result['success']:
            data = result['data']
            return json.dumps({
                'status': 'success',
                'progress_id': data.get('progress_id'),
                'user_id': data.get('user_id'),
                'contents_id': data.get('contents_id'),
                'textbook_id': data.get('textbook_id'),
                'seq_number': data.get('seq_number'),
                'section_title': data.get('section_title'),
                'progress': data.get('progress'),
                'previous_progress': data.get('previous_progress'),
                'message': data.get('message')
            }, ensure_ascii=False, indent=2)
        else:
            return json.dumps({
                'status': 'error',
                'message': result['error']
            }, ensure_ascii=False, indent=2)
            
    except Exception as e:
        return json.dumps({
            'status': 'error',
            'message': f'Failed to save progress: {str(e)}'
        }, ensure_ascii=False, indent=2)

def llm_get_problems_by_topic(subject_id: int, topic_id: int, user_id: int) -> str:
    """
    LLM Function: Get problems by topic ID for a specific subject (no grade filtering)
    
    Args:
        subject_id: ID of the subject
        topic_id: ID of the topic to search for
        user_id: ID of the user
        
    Returns:
        JSON string with problems for the specified topic or error message
    """
    try:
        # Convert subject_id to int if it's a string (from Gemini API)
        if isinstance(subject_id, str):
            try:
                subject_id = int(subject_id)
            except ValueError:
                return json.dumps({
                    'status': 'error',
                    'message': 'subject_id must be a valid integer'
                }, ensure_ascii=False, indent=2)
        
        if not isinstance(subject_id, int):
            return json.dumps({
                'status': 'error',
                'message': 'subject_id must be an integer'
            }, ensure_ascii=False, indent=2)
        
        # Validate topic_id is an integer
        if not isinstance(topic_id, int):
            # Try to convert topic_id to int if it's a string (from Gemini API)
            if isinstance(topic_id, str):
                try:
                    topic_id = int(topic_id)
                except ValueError:
                    return json.dumps({
                        'status': 'error',
                        'message': 'topic_id must be a valid integer'
                    }, ensure_ascii=False, indent=2)
            else:
                return json.dumps({
                    'status': 'error',
                    'message': 'topic_id must be an integer'
                }, ensure_ascii=False, indent=2)
        
        client = get_client()
        result = client.get_problems_by_topic(subject_id, topic_id, user_id)
        
        if result['success']:
            data = result['data']
            return json.dumps({
                'status': 'success',
                'subject_id': data.get('subject_id'),
                'subject_folder': data.get('subject_folder'),
                'topic_id': data.get('topic_id'),
                'topic_name': data.get('topic_name'),
                'total_problems': data.get('total_problems'),
                'problems': data.get('problems')
            }, ensure_ascii=False, indent=2)
        else:
            return json.dumps({
                'status': 'error',
                'message': result['error']
            }, ensure_ascii=False, indent=2)
            
    except Exception as e:
        return json.dumps({
            'status': 'error',
            'message': f'Failed to get problems by topic: {str(e)}'
        }, ensure_ascii=False, indent=2)

def llm_mark_problem_as_solved(user_id: int, problem_id: int) -> str:
    """
    LLM Function: Mark a problem as solved for a specific user
    
    Args:
        problem_id: ID of the problem to mark as solved
        user_id: ID of the user
        
    Returns:
        JSON string with success/error status and details
    """
    try:
        # Convert problem_id to int if it's a string
        if isinstance(problem_id, str):
            try:
                problem_id = int(problem_id)
            except ValueError:
                return json.dumps({
                    'status': 'error',
                    'message': 'problem_id must be a valid integer'
                }, ensure_ascii=False, indent=2)

        if not isinstance(problem_id, int):
            return json.dumps({
                'status': 'error',
                'message': 'problem_id must be an integer'
            }, ensure_ascii=False, indent=2)
            
        client = get_client()
        result = client.mark_problem_as_solved(user_id, problem_id)
        
        if result['success']:
            return json.dumps({
                'status': 'success',
                'message': result['data'].get('message')
            }, ensure_ascii=False, indent=2)
        else:
            return json.dumps({
                'status': 'error',
                'message': result['error']
            }, ensure_ascii=False, indent=2)

    except Exception as e:
        return json.dumps({
            'status': 'error',
            'message': f'Failed to mark problem as solved: {str(e)}'
        }, ensure_ascii=False, indent=2)


# === FUNCTION DEFINITIONS FOR LLM ===

def get_llm_function_definitions():
    """
    Get function definitions for LLM function calling
    
    Returns:
        List of function definitions in OpenAI format
    """
    return [
        {
            "name": "llm_get_question_from_db",
            "description": "Получить конкретный вопрос из базы данных по предмету, классу и номеру вопроса. ВАЖНО: Начинай всегда с question_order=1 - это покажет общее количество вопросов в тесте и первый вопрос. Затем получай следующие вопросы используя next_question_order из ответа, пока is_last_question не станет true. Задавай вопросы пользователю по одному и запрашивай следующий только после ответа пользователя. В ответе ты получишь: current_question (сам вопрос и ответ), progress (прогресс X/Y и количество оставшихся), next_action (инструкции для следующего шага).",
            "parameters": {
                "type": "object",
                "properties": {
                    "subject_id": {
                        "type": "integer",
                        "description": "ID предмета (1-физика, 2-алгебра, 3-геометрия, 4-химия, 5-биология, 6-история, 7-русский, 8-обществознание, 9-литература, 10-география, 11-английский, 12-информатика)"
                    },
                    "grade": {
                        "type": "string",
                        "description": "Класс как число: 7, 8, 9, 10, 11",
                        "enum": ["7", "8", "9", "10", "11"]
                    },
                    "question_order": {
                        "type": "integer",
                        "description": "Номер вопроса в базе данных. Начинай с 1 - это покажет общее количество вопросов и первый вопрос. Затем используй next_question_order из предыдущего ответа",
                        "minimum": 1
                    }
                },
                "required": ["subject_id", "grade", "question_order"]
            }
        },
        {
            "name": "llm_get_problems_by_topic",
            "description": "Получить задачи по теме для конкретного предмета БЕЗ учета класса. Ищи задачи в таблице tu_problems_topic по ID темы, затем получай все задачи из tu_problem с найденным topic_id. Задачи сортируются по порядку (q_order).",
            "parameters": {
                "type": "object",
                "properties": {
                    "subject_id": {
                        "type": "integer",
                        "description": "ID предмета (1-физика, 2-алгебра, 3-геометрия, 4-химия, 5-биология, 6-история, 7-русский, 8-обществознание, 9-литература, 10-география, 11-английский, 12-информатика)"
                    },
                    "topic_id": {
                        "type": "integer",
                        "description": "ID темы для поиска",
                        "minimum": 1
                    },
                    "user_id": {
                        "type": "integer",
                        "description": "ID пользователя"
                    }
                },
                "required": ["subject_id", "topic_id", "user_id"]
            }
        },
        {
            "name": "llm_list_available_topics",
            "description": "Получить список доступных тем для конкретного предмета (без фильтрации по классу). Возвращает ID и название темы.",
            "parameters": {
                "type": "object",
                "properties": {
                    "subject_id": {
                        "type": "integer",
                        "description": "ID предмета (1-физика, 2-алгебра, 3-геометрия, 4-химия, 5-биология, 6-история, 7-русский, 8-обществознание, 9-литература, 10-география, 11-английский, 12-информатика)"
                    }
                },
                "required": ["subject_id"]
            }
        },
        {
            "name": "llm_mark_problem_as_solved",
            "description": "Пометить задачу как решенную для конкретного пользователя. Используй эту функцию после того, как показал задачу ученику.",
            "parameters": {
                "type": "object",
                "properties": {
                    "problem_id": {
                        "type": "integer",
                        "description": "ID задачи, которую нужно пометить как решенную"
                    },
                    "user_id": {
                        "type": "integer",
                        "description": "ID пользователя (необязательный, будет взят из сессии)"
                    }
                },
                "required": ["problem_id", "user_id"]
            }
        },
        {
            "name": "llm_save_paragraph_progress",
            "description": "Сохранить прогресс студента по конкретному параграфу. Укажи ID пользователя, ID учебника, номер параграфа (seq_number) и текст прогресса. В ответе ты получишь: progress_id, user_id, contents_id, textbook_id, seq_number, section_title, progress, previous_progress, message.",
            "parameters": {
                "type": "object",
                "properties": {
                    "user_id": {
                        "type": "integer",
                        "description": "ID пользователя"
                    },
                    "textbook_id": {
                        "type": "integer",
                        "description": "ID учебника"
                    },
                    "seq_number": {
                        "type": "string",
                        "description": "Номер параграфа из таблицы contents (например, \"1.1\", \"2.3\", и т.д.)"
                    },
                    "progress": {
                        "type": "string",
                        "description": "Текст, описывающий, что студент изучил",
                        "minLength": 1,
                        "maxLength": 1000
                    }
                },
                "required": ["user_id", "textbook_id", "seq_number", "progress"]
            }
        }
    ]

def cleanup_llm_client():
    """Cleanup global client when shutting down"""
    global _global_client
    if _global_client:
        _global_client._stop_server()
        _global_client = None


if __name__ == '__main__':
    import sys