"""
MCP Client implementation for Spark optimization.
"""
import asyncio
import json
from typing import Dict, List, Any, Optional

class SparkMCPClient:
    """MCP client for interacting with Spark optimization server."""
    
    def __init__(self, host: str = "localhost", port: int = 8080):
        self.host = host
        self.port = port
        self._reader: Optional[asyncio.StreamReader] = None
        self._writer: Optional[asyncio.StreamWriter] = None

    async def connect(self):
        """Connect to the MCP server."""
        self._reader, self._writer = await asyncio.open_connection(
            self.host, self.port
        )

    async def close(self):
        """Close the connection to the MCP server."""
        if self._writer:
            self._writer.close()
            await self._writer.wait_closed()

    async def list_tools(self) -> List[Dict[str, Any]]:
        """List available tools from the MCP server."""
        response = await self._send_request({
            "method": "tools/list"
        })
        return response["tools"]

    async def optimize_spark_code(self, code: str, optimization_level: str = "advanced", save_to_file: bool = True) -> str:
        """Optimize Spark code using the MCP server.
        
        Args:
            code: The Spark code to optimize
            optimization_level: Level of optimization ("basic" or "advanced")
            save_to_file: If True, saves optimized code to output/optimized_spark_example.py
            
        Returns:
            The optimized code as a string
        """
        response = await self._send_request({
            "method": "tools/call",
            "params": {
                "name": "optimize_spark_code",
                "arguments": {
                    "code": code,
                    "optimization_level": optimization_level
                }
            }
        })
        
        optimized_code = response["content"][0]["text"]
        
        if save_to_file:
            # Extract just the code block from the response
            code_block = ""
            in_code_block = False
            for line in optimized_code.split("\n"):
                if line.startswith("```python"):
                    in_code_block = True
                    continue
                elif line.startswith("```"):
                    in_code_block = False
                    continue
                elif in_code_block:
                    code_block += line + "\n"
            
            # Save to file
            import os
            output_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "output")
            os.makedirs(output_dir, exist_ok=True)
            output_file = os.path.join(output_dir, "optimized_spark_example.py")
            
            # Get performance analysis for optimization comments
            perf_analysis = await self.analyze_performance(code, optimized_code)
            
            # Extract optimization details from performance analysis
            optimization_comments = []
            if isinstance(perf_analysis, dict) and 'text' in perf_analysis:
                analysis_text = perf_analysis['text']
                # Extract key points from analysis
                for line in analysis_text.split('\n'):
                    if any(key in line.lower() for key in ['improvement', 'optimization', 'performance', 'tuning']):
                        if line.strip().startswith('-') or line.strip().startswith('•'):
                            optimization_comments.append(line.strip())
            
            # Create detailed docstring
            docstring = '"""\nOptimized PySpark code generated by Spark MCP.\n\nOptimization Details:\n-------------------\n'
            for i, comment in enumerate(optimization_comments, 1):
                docstring += f'{i}. {comment.lstrip("-•").strip()}\n'
            docstring += '"""\n'
            # Extract key optimizations from performance analysis
            optimizations = {}
            if isinstance(perf_analysis, dict) and 'text' in perf_analysis:
                analysis_text = perf_analysis['text']
                current_section = ''
                current_details = []
                
                for line in analysis_text.split('\n'):
                    line = line.strip()
                    if not line:
                        continue
                        
                    # Start of a new section
                    if line.endswith(':'):
                        if current_section and current_details:
                            optimizations[current_section] = current_details
                        current_section = line[:-1]
                        current_details = []
                    # Detail line
                    elif line.startswith('-') or line.startswith('•'):
                        current_details.append(line.lstrip('-• ').strip())
                        
                if current_section and current_details:
                    optimizations[current_section] = current_details
            
            # Add inline comments to the code
            code_lines = code_block.split('\n')
            commented_code = []
            
            # Process each line
            for line in code_lines:
                commented_code.append(line)
                
                # Look for optimization opportunities in the line
                line_lower = line.strip().lower()
                
                # Find relevant optimizations for this line
                for section, details in optimizations.items():
                    if any(keyword in line_lower for keyword in section.lower().split()):
                        commented_code.append(f'# Performance Optimization: {section}')
                        for detail in details:
                            commented_code.append(f'# - {detail}')
                        commented_code.append('')
            
            # Write to file
            with open(output_file, "w") as f:
                f.write(docstring)
                f.write('\n'.join(commented_code))
            
            print(f"\nOptimized code saved to: {output_file}")
        
        return optimized_code

    async def analyze_performance(self, original_code: str, optimized_code: str, save_to_file: bool = True) -> Dict[str, Any]:
        """Analyze performance between original and optimized code.
        
        Args:
            original_code: The original Spark code
            optimized_code: The optimized Spark code
            save_to_file: If True, saves analysis to output/performance_analysis.md
            
        Returns:
            Performance analysis results
        """
        response = await self._send_request({
            "method": "tools/call",
            "params": {
                "name": "analyze_performance",
                "arguments": {
                    "original_code": original_code,
                    "optimized_code": optimized_code
                }
            }
        })
        analysis = response["content"][0]
        
        if save_to_file and isinstance(analysis, dict) and 'text' in analysis:
            import os
            output_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "output")
            os.makedirs(output_dir, exist_ok=True)
            
            # Format the analysis as markdown
            markdown_content = "# Spark Code Performance Analysis\n\n"
            markdown_content += analysis['text'].replace('\n', '\n\n')
            
            # Save to file
            output_file = os.path.join(output_dir, "performance_analysis.md")
            with open(output_file, "w") as f:
                f.write(markdown_content)
            
            print(f"\nPerformance analysis saved to: {output_file}")
            
        return analysis

    async def _send_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """Send a request to the MCP server and get the response."""
        if not self._writer or not self._reader:
            raise RuntimeError("Not connected to MCP server")

        self._writer.write(json.dumps(request).encode() + b"\n")
        await self._writer.drain()

        data = await self._reader.readline()
        if not data:
            raise ConnectionError("Connection closed by server")

        return json.loads(data)
