#!/usr/bin/env python3
"""
Connections MCP Server
"""

import json
import random
import asyncio
import sys
import uuid
from pathlib import Path
from typing import List, Dict, Tuple, Optional, Any
from datetime import datetime

from mcp.server import Server, NotificationOptions
from mcp.server.models import InitializationOptions
import mcp.server.stdio
import mcp.types as types


class ConnectionsPuzzle:
    def __init__(self, filename: str = "connections.txt"):
        self.categories = {}
        self.puzzles = {}  # Store puzzles by ID
        self.load_categories(filename)
    
    def load_categories(self, filename: str):
        """Load categories and words from the file"""
        # Try to find the file in multiple locations
        possible_paths = [
            Path(__file__).parent / filename,  # Same directory as script
            Path(sys.prefix) / filename,  # In the Python installation
            Path.cwd() / filename,  # Current working directory
        ]
        
        # For installed packages, check if file is in site-packages
        try:
            import site
            for site_dir in site.getsitepackages():
                possible_paths.append(Path(site_dir) / filename)
        except:
            pass
        
        filepath = None
        for path in possible_paths:
            if path.exists():
                filepath = path
                break
        
        if filepath is None:
            # Try to read from package data if this is an installed package
            try:
                import importlib.resources as pkg_resources
                # For Python 3.9+
                if hasattr(pkg_resources, 'files'):
                    data = pkg_resources.files('connections_mcp').joinpath(filename).read_text()
                    lines = data.splitlines(keepends=True)
                else:
                    # For Python 3.7-3.8
                    with pkg_resources.open_text('connections_mcp', filename) as f:
                        lines = f.readlines()
            except:
                # Last resort: try to read the embedded data
                lines = self._get_embedded_data().splitlines(keepends=True)
        else:
            with open(filepath, 'r') as f:
                lines = f.readlines()
        
        # Skip header lines and find where the data starts
        start_index = 0
        for i, line in enumerate(lines):
            if line.strip() == "group, words:":
                start_index = i + 1
                break
        
        # Parse categories and words
        for i in range(start_index, len(lines)):
            line = lines[i].strip()
            if line and ',' in line:
                parts = line.split(',', 1)
                if len(parts) == 2:
                    category = parts[0].strip()
                    words = parts[1].strip().split()
                    if len(words) == 4:
                        self.categories[category] = words
    
    def _get_embedded_data(self) -> str:
        """Return embedded data as a fallback"""
        return """Below is some raw data for Connections daily game.
The game introduces 16 words, from which one should identify 4 different groups of words, each one of them should include distinct 4 words out of the 16 provided.

group, words:
MUSICAL INSTRUMENTS,BASS BASSOON HARP RECORDER
PLANT GROWTHS,BLOOM BUD SHOOT SPROUT
BRING UP,FOSTER NURSE RAISE REAR
SOLAR EMANATIONS,CORONA FLARE LIGHT RADIATION
HOW FAST SOMETHING IS GOING,CLIP PACE RATE SPEED
N.B.A. TEAM PLAYER,CLIPPER PACER ROCKET SPUR
FAIL TO ATTEND,CUT DITCH MISS SKIP
DECISIVE VICTORY,BLOWOUT ROMP ROUT SWEEP
KUDOS,CREDIT GLORY PRAISE PROPS
BECOME AWARE OF,DISCOVER FIND LEARN REALIZE
PURSUE,CHASE TAIL TRACK SHADOW
SEEN IN A PASSPORT,NAME PHOTO STAMP VISA
JUMPING ANIMALS,CRICKET FROG HARE KANGAROO
APPLY PRESSURE TO,CRUSH MASH PRESS SQUASH
OLYMPIC SPORTS,BREAKING HOCKEY SKELETON TRAMPOLINE
THINGS YOU CAN SET,MOOD RECORD TABLE VOLLEYBALL"""
    
    def create_puzzle(self, num_categories: int = 4) -> Tuple[str, List[str], Dict[str, List[str]]]:
        """Create a puzzle by selecting random categories and shuffling words"""
        if num_categories > len(self.categories):
            raise ValueError(f"Cannot create puzzle with {num_categories} categories, only {len(self.categories)} available")
        
        # Generate unique puzzle ID
        puzzle_id = str(uuid.uuid4())
        
        # Select random categories
        selected_categories = random.sample(list(self.categories.keys()), num_categories)
        
        # Collect all words and shuffle
        all_words = []
        solution = {}
        
        for category in selected_categories:
            words = self.categories[category]
            all_words.extend(words)
            solution[category] = words.copy()
        
        random.shuffle(all_words)
        
        # Store puzzle with ID
        self.puzzles[puzzle_id] = {
            "words": all_words,
            "solution": solution,
            "created_at": datetime.now().isoformat()
        }
        
        return puzzle_id, all_words, solution
    
    def test_solution(self, puzzle_id: str, proposed_groups: List[List[str]]) -> Dict[str, Any]:
        """Test a proposed solution against a specific puzzle"""
        if puzzle_id not in self.puzzles:
            return {
                "error": f"Puzzle with ID {puzzle_id} not found.",
                "correct_groups": [],
                "incorrect_groups": [],
                "is_complete": False
            }
        
        puzzle = self.puzzles[puzzle_id]
        solution = puzzle["solution"]
        
        result = {
            "correct_groups": [],
            "incorrect_groups": [],
            "is_complete": False,
            "message": ""
        }
        
        # Check each proposed group
        for i, group in enumerate(proposed_groups):
            if len(group) != 4:
                result["incorrect_groups"].append({
                    "group_index": i + 1,
                    "words": group,
                    "issue": f"Group has {len(group)} words, should have 4"
                })
                continue
            
            # Check if this group matches any category
            group_set = set([word.upper() for word in group])
            matched = False
            
            for category, correct_words in solution.items():
                if group_set == set(correct_words):
                    result["correct_groups"].append({
                        "group_index": i + 1,
                        "words": group,
                        "category": category
                    })
                    matched = True
                    break
            
            if not matched:
                result["incorrect_groups"].append({
                    "group_index": i + 1,
                    "words": group,
                    "issue": "Words don't form a valid category"
                })
        
        # Check if solution is complete
        num_correct = len(result["correct_groups"])
        num_total = len(solution)
        
        if num_correct == num_total and len(proposed_groups) == num_total:
            result["is_complete"] = True
            result["message"] = "Congratulations! All groups are correct!"
        elif num_correct > 0:
            result["message"] = f"Partial solution: {num_correct} out of {num_total} groups are correct."
        else:
            result["message"] = "No correct groups found."
        
        return result
    
    def get_puzzle(self, puzzle_id: str) -> Optional[Dict[str, Any]]:
        """Get a specific puzzle by ID"""
        if puzzle_id not in self.puzzles:
            return None
        return {
            "puzzle_id": puzzle_id,
            "words": self.puzzles[puzzle_id]["words"],
            "created_at": self.puzzles[puzzle_id]["created_at"]
        }
    
    def list_puzzles(self) -> List[Dict[str, str]]:
        """List all available puzzles"""
        return [
            {
                "puzzle_id": pid,
                "created_at": puzzle["created_at"]
            }
            for pid, puzzle in self.puzzles.items()
        ]
    
    def get_hint(self, puzzle_id: str, category_index: int = 0) -> Optional[str]:
        """Get a hint for a specific puzzle"""
        if puzzle_id not in self.puzzles:
            return None
        
        solution = self.puzzles[puzzle_id]["solution"]
        categories = list(solution.keys())
        if category_index >= len(categories):
            return None
        
        category = categories[category_index]
        return f"One category is: {category}"


# Create the MCP server
server = Server("connections-mcp")
puzzle_game = None  # Will be initialized in main


@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    """List available tools"""
    return [
        types.Tool(
            name="create_puzzle",
            description="Create a new Connections puzzle with shuffled words from 4 categories",
            inputSchema={
                "type": "object",
                "properties": {
                    "num_categories": {
                        "type": "integer",
                        "description": "Number of categories (default: 4)",
                        "default": 4
                    }
                },
                "required": []
            },
        ),
        types.Tool(
            name="test_solution",
            description="Test a proposed solution for a specific puzzle",
            inputSchema={
                "type": "object",
                "properties": {
                    "puzzle_id": {
                        "type": "string",
                        "description": "The ID of the puzzle to test"
                    },
                    "groups": {
                        "type": "array",
                        "description": "Array of 4 groups, each containing 4 words",
                        "items": {
                            "type": "array",
                            "items": {"type": "string"}
                        }
                    }
                },
                "required": ["puzzle_id", "groups"]
            },
        ),
        types.Tool(
            name="get_puzzle",
            description="Get a specific puzzle by ID",
            inputSchema={
                "type": "object",
                "properties": {
                    "puzzle_id": {
                        "type": "string",
                        "description": "The ID of the puzzle to retrieve"
                    }
                },
                "required": ["puzzle_id"]
            },
        ),
        types.Tool(
            name="list_puzzles",
            description="List all available puzzles",
            inputSchema={
                "type": "object",
                "properties": {},
                "required": []
            },
        ),
        types.Tool(
            name="get_hint",
            description="Get a hint for a specific puzzle",
            inputSchema={
                "type": "object",
                "properties": {
                    "puzzle_id": {
                        "type": "string",
                        "description": "The ID of the puzzle"
                    },
                    "category_index": {
                        "type": "integer",
                        "description": "Which category hint to reveal (0-3)",
                        "default": 0
                    }
                },
                "required": ["puzzle_id"]
            },
        ),
    ]


@server.call_tool()
async def handle_call_tool(
    name: str, arguments: dict | None
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """Handle tool calls"""
    
    if name == "create_puzzle":
        num_categories = arguments.get("num_categories", 4) if arguments else 4
        try:
            puzzle_id, words, solution = puzzle_game.create_puzzle(num_categories)
            
            # Format puzzle as 4x4 grid
            grid = []
            for i in range(0, len(words), 4):
                grid.append(words[i:i+4])
            
            response = {
                "puzzle_id": puzzle_id,
                "grid": grid,
                "words": words,
                "message": "New puzzle created! Find 4 groups of 4 related words."
            }
            
            return [types.TextContent(
                type="text",
                text=json.dumps(response, indent=2)
            )]
        except Exception as e:
            return [types.TextContent(
                type="text",
                text=json.dumps({"error": str(e)})
            )]
    
    elif name == "test_solution":
        if not arguments or "groups" not in arguments or "puzzle_id" not in arguments:
            return [types.TextContent(
                type="text",
                text=json.dumps({"error": "Both puzzle_id and groups parameters are required"})
            )]
        
        puzzle_id = arguments["puzzle_id"]
        groups = arguments["groups"]
        result = puzzle_game.test_solution(puzzle_id, groups)
        
        return [types.TextContent(
            type="text",
            text=json.dumps(result, indent=2)
        )]
    
    elif name == "get_puzzle":
        if not arguments or "puzzle_id" not in arguments:
            return [types.TextContent(
                type="text",
                text=json.dumps({"error": "puzzle_id parameter is required"})
            )]
        
        puzzle_id = arguments["puzzle_id"]
        puzzle_data = puzzle_game.get_puzzle(puzzle_id)
        
        if not puzzle_data:
            return [types.TextContent(
                type="text",
                text=json.dumps({"error": f"Puzzle with ID {puzzle_id} not found"})
            )]
        
        # Format as grid
        words = puzzle_data["words"]
        grid = []
        for i in range(0, len(words), 4):
            grid.append(words[i:i+4])
        
        response = {
            "puzzle_id": puzzle_data["puzzle_id"],
            "grid": grid,
            "words": words,
            "created_at": puzzle_data["created_at"]
        }
        
        return [types.TextContent(
            type="text",
            text=json.dumps(response, indent=2)
        )]
    
    elif name == "list_puzzles":
        puzzles = puzzle_game.list_puzzles()
        return [types.TextContent(
            type="text",
            text=json.dumps({"puzzles": puzzles}, indent=2)
        )]
    
    elif name == "get_hint":
        if not arguments or "puzzle_id" not in arguments:
            return [types.TextContent(
                type="text",
                text=json.dumps({"error": "puzzle_id parameter is required"})
            )]
        
        puzzle_id = arguments["puzzle_id"]
        category_index = arguments.get("category_index", 0)
        
        hint = puzzle_game.get_hint(puzzle_id, category_index)
        if hint:
            return [types.TextContent(
                type="text",
                text=json.dumps({"hint": hint})
            )]
        else:
            return [types.TextContent(
                type="text",
                text=json.dumps({"error": f"No hint available for puzzle {puzzle_id}"})
            )]
    
    else:
        return [types.TextContent(
            type="text",
            text=json.dumps({"error": f"Unknown tool: {name}"})
        )]


async def amain():
    """Run the MCP server"""
    global puzzle_game
    try:
        puzzle_game = ConnectionsPuzzle()
    except Exception as e:
        print(f"Error initializing ConnectionsPuzzle: {e}", file=sys.stderr)
        # Continue with embedded data
        puzzle_game = ConnectionsPuzzle()
    
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="connections-mcp",
                server_version="0.2.1",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )


def main():
    """Entry point for the package"""
    asyncio.run(amain())


if __name__ == "__main__":
    main()