"""
BML MCP Server - Build-Measure-Learn GitHub project management for AI agents
"""
import json
from enum import Enum
from typing import Sequence

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent, ImageContent, EmbeddedResource
from mcp.shared.exceptions import McpError
from pydantic import BaseModel


class BMLTools(str, Enum):
    """Available BML tools for GitHub project management"""
    LIST_ISSUES = "list_issues"
    GET_ISSUE = "get_issue"
    CREATE_ISSUE = "create_issue"
    MOVE_ISSUE_ABOVE = "move_issue_above"
    MOVE_ISSUE_BELOW = "move_issue_below"
    MOVE_ISSUE_BETWEEN = "move_issue_between"
    SET_ISSUE_STATUS = "set_issue_status"
    VIEW_KANBAN = "view_kanban"
    GET_KANBAN_LANE = "get_kanban_lane"


# Import real BML functions from heaven-bml-system package
try:
    from heaven_bml.tree_kanban import (
        get_all_prioritized_issues,
        move_issue_above,
        move_issue_below, 
        move_issue_between
    )
    from heaven_bml.github_kanban import (
        construct_kanban_from_labels,
        view_lane
    )
    print("✅ Successfully imported real BML functions")
    USING_REAL_FUNCTIONS = True
except ImportError as e:
    print(f"⚠️ Could not import real BML functions ({e}), using stubs")
    USING_REAL_FUNCTIONS = False
    
    # Fallback stub implementations
    def get_all_prioritized_issues(repo: str) -> list:
        return [
            {"number": 1, "title": "Sample Issue 1", "labels": ["status-backlog", "priority-1"]},
            {"number": 2, "title": "Sample Issue 2", "labels": ["status-plan", "priority-2"]},
        ]
    
    def move_issue_above(issue_id: str, target_issue_id: str, repo: str) -> str:
        return f"STUB: Moved issue {issue_id} above issue {target_issue_id} in {repo}"
    
    def move_issue_below(issue_id: str, target_issue_id: str, repo: str) -> str:
        return f"STUB: Moved issue {issue_id} below issue {target_issue_id} in {repo}"
    
    def move_issue_between(issue_id: str, above_issue_id: str, below_issue_id: str, repo: str) -> str:
        return f"STUB: Moved issue {issue_id} between {above_issue_id} and {below_issue_id} in {repo}"
    
    def construct_kanban_from_labels(repo: str) -> dict:
        return {
            "backlog": [{"number": 1, "title": "Sample Issue 1"}],
            "plan": [{"number": 2, "title": "Sample Issue 2"}],
            "build": [], "measure": [], "learn": [], "blocked": [], "archived": []
        }
    
    def view_lane(repo: str, status: str) -> list:
        return [{"number": 1, "title": f"Sample issue in {status}"}]

# Real implementation of set_issue_status using GitHub CLI
def set_issue_status(issue_id: str, status: str, repo: str) -> str:
    """Set issue status by updating GitHub labels"""
    if USING_REAL_FUNCTIONS:
        # Use real BML function
        from heaven_bml.tree_kanban import set_issue_status as bml_set_status
        success = bml_set_status(repo, int(issue_id), status)
        if success:
            return f"✅ Set issue #{issue_id} to status '{status}' in {repo}"
        else:
            return f"❌ Failed to set issue #{issue_id} to status '{status}' in {repo}"
    else:
        return f"STUB: Set issue {issue_id} to status {status} in {repo}"


class BMLServer:
    """
    HEAVEN Build-Measure-Learn Server with Tree Notation Priorities
    
    This server provides AI agents with sophisticated GitHub project management using:
    
    🌳 TREE NOTATION PRIORITIES:
    - Hierarchical priorities: 1 > 1.1 > 1.2 > 1.2.1 > 1.2.2 > 1.3 > 2
    - Infinite nesting allows precise task organization
    - Natural sorting: priority-1.2.3 comes before priority-1.2.4
    - Create sub-tasks by moving below parent (1.2 → 1.2.1, 1.2.2...)
    
    📋 BUILD-MEASURE-LEARN WORKFLOW:
    - backlog: Ideas and requests waiting for prioritization
    - plan: Prioritized and designed tasks ready for development  
    - build: Currently in development
    - measure: Testing, validation, getting feedback
    - learn: Analysis, insights, retrospective
    - blocked: Stuck at any stage (with reason)
    - archived: Completed work
    
    🔧 OPERATIONS:
    Use move_issue_above/below/between to create tree hierarchies.
    Use set_issue_status to progress through BML workflow.
    Use view_kanban to see complete project state with priorities.
    """
    
    def __init__(self, default_repo: str = "sancovp/heaven-base"):
        self.default_repo = default_repo
    
    def list_issues(self, repo: str = None) -> dict:
        """Get all prioritized issues for a repository"""
        repo = repo or self.default_repo
        issues = get_all_prioritized_issues(repo)
        return {"repo": repo, "issues": issues, "total": len(issues)}
    
    def get_issue(self, issue_id: str, repo: str = None) -> dict:
        """Get details for a specific issue"""
        repo = repo or self.default_repo
        
        if USING_REAL_FUNCTIONS:
            # Use GitHub CLI to get issue details
            import subprocess
            import json
            
            try:
                cmd = f'gh issue view {issue_id} --repo {repo} --json number,title,body,labels,state'
                result = subprocess.run(cmd, shell=True, capture_output=True, text=True, check=True)
                
                # Debug JSON parsing
                try:
                    issue_data = json.loads(result.stdout)
                except json.JSONDecodeError as e:
                    print(f"JSON decode error in get_issue: {e}")
                    print(f"Command output: {result.stdout[:500]}...")
                    return {"repo": repo, "issue_id": issue_id, "error": "JSON parse error"}
                
                # Extract status from labels
                status = 'backlog'  # default
                priority = 'none'
                for label in issue_data.get('labels', []):
                    if label['name'].startswith('status-'):
                        status = label['name'][7:]  # Remove 'status-' prefix
                    elif label['name'].startswith('priority-'):
                        priority = label['name'][9:]  # Remove 'priority-' prefix
                
                return {
                    "repo": repo,
                    "issue_id": issue_id,
                    "number": issue_data['number'],
                    "title": issue_data['title'],
                    "body": issue_data['body'],
                    "state": issue_data['state'],
                    "status": status,
                    "priority": priority
                }
            except subprocess.CalledProcessError:
                return {"repo": repo, "issue_id": issue_id, "error": "Issue not found"}
        else:
            return {"repo": repo, "issue_id": issue_id, "status": "placeholder"}
    
    def create_issue(self, title: str, body: str = "", labels: list = None, repo: str = None) -> dict:
        """Create a new issue"""
        repo = repo or self.default_repo
        labels = labels or ["status-backlog", "priority-medium"]
        
        if USING_REAL_FUNCTIONS:
            # Use real BML function to create GitHub issue
            from heaven_bml.github_kanban import create_github_issue_with_status
            issue_id = create_github_issue_with_status(repo, title, body, "backlog")  # No priority by default
            return {"repo": repo, "title": title, "body": body, "issue_id": issue_id, "created": True}
        else:
            # Stub implementation
            return {"repo": repo, "title": title, "body": body, "labels": labels, "created": True}
    
    def move_issue_above(self, issue_id: str, target_issue_id: str, repo: str = None) -> str:
        """Move issue above target issue in priority"""
        repo = repo or self.default_repo
        return move_issue_above(issue_id, target_issue_id, repo)
    
    def move_issue_below(self, issue_id: str, target_issue_id: str, repo: str = None) -> str:
        """Move issue below target issue in priority"""
        repo = repo or self.default_repo
        return move_issue_below(issue_id, target_issue_id, repo)
    
    def move_issue_between(self, issue_id: str, above_issue_id: str, below_issue_id: str, repo: str = None) -> str:
        """Move issue between two other issues"""
        repo = repo or self.default_repo
        return move_issue_between(issue_id, above_issue_id, below_issue_id, repo)
    
    def set_issue_status(self, issue_id: str, status: str, repo: str = None) -> str:
        """Set issue status (backlog/plan/build/measure/learn/blocked/archived)"""
        repo = repo or self.default_repo
        return set_issue_status(issue_id, status, repo)
    
    def view_kanban(self, repo: str = None) -> dict:
        """Get complete kanban board view"""
        repo = repo or self.default_repo
        kanban = construct_kanban_from_labels(repo)
        return {
            "repo": repo, 
            "kanban": {
                "backlog": [{"number": issue.number, "title": issue.title} for issue in kanban.backlog],
                "plan": [{"number": issue.number, "title": issue.title} for issue in kanban.plan],
                "build": [{"number": issue.number, "title": issue.title} for issue in kanban.build],
                "measure": [{"number": issue.number, "title": issue.title} for issue in kanban.measure],
                "learn": [{"number": issue.number, "title": issue.title} for issue in kanban.learn],
                "blocked": [{"number": issue.number, "title": issue.title} for issue in kanban.blocked],
                "archived": [{"number": issue.number, "title": issue.title} for issue in kanban.archived]
            }
        }
    
    def get_kanban_lane(self, status: str, repo: str = None) -> dict:
        """Get issues in specific kanban lane"""
        repo = repo or self.default_repo
        issues = view_lane(repo, status)
        return {"repo": repo, "status": status, "issues": issues}


async def serve(default_repo: str = "sancovp/heaven-base") -> None:
    """Main MCP server function"""
    server = Server("heaven-bml")
    bml_server = BMLServer(default_repo)
    
    @server.list_tools()
    async def list_tools() -> list[Tool]:
        """List available BML tools"""
        return [
            Tool(
                name=BMLTools.LIST_ISSUES.value,
                description="List all issues with HEAVEN tree notation priorities. Tree notation uses hierarchical priorities like 1.2.3.4.5 where 1.2.3 comes before 1.2.4. This allows infinite nesting and precise ordering of priorities.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "repo": {
                            "type": "string",
                            "description": f"GitHub repository (default: {default_repo})",
                        }
                    },
                    "required": [],
                },
            ),
            Tool(
                name=BMLTools.GET_ISSUE.value,
                description="Get details for a specific issue",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "issue_id": {
                            "type": "string",
                            "description": "GitHub issue ID",
                        },
                        "repo": {
                            "type": "string", 
                            "description": f"GitHub repository (default: {default_repo})",
                        }
                    },
                    "required": ["issue_id"],
                },
            ),
            Tool(
                name=BMLTools.CREATE_ISSUE.value,
                description="Create a new GitHub issue",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "title": {
                            "type": "string",
                            "description": "Issue title",
                        },
                        "body": {
                            "type": "string",
                            "description": "Issue body/description",
                        },
                        "labels": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "Issue labels",
                        },
                        "repo": {
                            "type": "string",
                            "description": f"GitHub repository (default: {default_repo})",
                        }
                    },
                    "required": ["title"],
                },
            ),
            Tool(
                name=BMLTools.MOVE_ISSUE_ABOVE.value,
                description="Move issue above target issue using HEAVEN tree notation. Creates new priority like if target is priority-2.1, moved issue becomes priority-2.0.5. Tree notation allows infinite precision: 1 > 1.1 > 1.2 > 1.2.1 > 1.2.2 > 1.3. Use this to create sub-priorities and organize hierarchical task structures.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "issue_id": {
                            "type": "string",
                            "description": "Issue ID to move",
                        },
                        "target_issue_id": {
                            "type": "string",
                            "description": "Target issue ID to move above",
                        },
                        "repo": {
                            "type": "string",
                            "description": f"GitHub repository (default: {default_repo})",
                        }
                    },
                    "required": ["issue_id", "target_issue_id"],
                },
            ),
            Tool(
                name=BMLTools.MOVE_ISSUE_BELOW.value,
                description="Move issue below target issue using HEAVEN tree notation. Creates lower priority like if target is priority-2.1, moved issue becomes priority-2.1.5. Perfect for creating sub-tasks under parent issues.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "issue_id": {
                            "type": "string",
                            "description": "Issue ID to move",
                        },
                        "target_issue_id": {
                            "type": "string",
                            "description": "Target issue ID to move below",
                        },
                        "repo": {
                            "type": "string",
                            "description": f"GitHub repository (default: {default_repo})",
                        }
                    },
                    "required": ["issue_id", "target_issue_id"],
                },
            ),
            Tool(
                name=BMLTools.MOVE_ISSUE_BETWEEN.value,
                description="Move issue between two other issues using HEAVEN tree notation. Precisely positions issue by creating new priority between existing ones. Example: between priority-1.2 and priority-1.3, creates priority-1.2.5.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "issue_id": {
                            "type": "string",
                            "description": "Issue ID to move",
                        },
                        "above_issue_id": {
                            "type": "string",
                            "description": "Issue ID to be above",
                        },
                        "below_issue_id": {
                            "type": "string", 
                            "description": "Issue ID to be below",
                        },
                        "repo": {
                            "type": "string",
                            "description": f"GitHub repository (default: {default_repo})",
                        }
                    },
                    "required": ["issue_id", "above_issue_id", "below_issue_id"],
                },
            ),
            Tool(
                name=BMLTools.SET_ISSUE_STATUS.value,
                description="Set issue status in Build-Measure-Learn workflow: backlog (ideas/requests) → plan (prioritized/designed) → build (in development) → measure (testing/validation) → learn (insights/analysis) → archived (completed). Also supports blocked status for any stage.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "issue_id": {
                            "type": "string",
                            "description": "Issue ID to update",
                        },
                        "status": {
                            "type": "string",
                            "enum": ["backlog", "plan", "build", "measure", "learn", "blocked", "archived"],
                            "description": "New status for the issue",
                        },
                        "repo": {
                            "type": "string",
                            "description": f"GitHub repository (default: {default_repo})",
                        }
                    },
                    "required": ["issue_id", "status"],
                },
            ),
            Tool(
                name=BMLTools.VIEW_KANBAN.value,
                description="Get complete Build-Measure-Learn kanban board with tree notation priorities. Shows all lanes: backlog → plan → build → measure → learn → archived, plus blocked items. Issues sorted by HEAVEN tree notation (1 > 1.1 > 1.2 > 1.2.1).",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "repo": {
                            "type": "string",
                            "description": f"GitHub repository (default: {default_repo})",
                        }
                    },
                    "required": [],
                },
            ),
            Tool(
                name=BMLTools.GET_KANBAN_LANE.value,
                description="Get issues in a specific Build-Measure-Learn lane with tree notation priorities. Each lane represents a workflow stage: backlog (gathering), plan (prioritizing), build (developing), measure (testing), learn (analyzing), blocked (stuck), archived (done).",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "status": {
                            "type": "string",
                            "enum": ["backlog", "plan", "build", "measure", "learn", "blocked", "archived"],
                            "description": "Kanban lane/status to view",
                        },
                        "repo": {
                            "type": "string",
                            "description": f"GitHub repository (default: {default_repo})",
                        }
                    },
                    "required": ["status"],
                },
            ),
        ]
    
    @server.call_tool()
    async def call_tool(
        name: str, arguments: dict
    ) -> Sequence[TextContent | ImageContent | EmbeddedResource]:
        """Handle tool calls for BML operations"""
        try:
            match name:
                case BMLTools.LIST_ISSUES.value:
                    result = bml_server.list_issues(arguments.get("repo"))
                
                case BMLTools.GET_ISSUE.value:
                    issue_id = arguments.get("issue_id")
                    if not issue_id:
                        raise ValueError("Missing required argument: issue_id")
                    result = bml_server.get_issue(issue_id, arguments.get("repo"))
                
                case BMLTools.CREATE_ISSUE.value:
                    title = arguments.get("title")
                    if not title:
                        raise ValueError("Missing required argument: title")
                    result = bml_server.create_issue(
                        title,
                        arguments.get("body", ""),
                        arguments.get("labels"),
                        arguments.get("repo")
                    )
                
                case BMLTools.MOVE_ISSUE_ABOVE.value:
                    issue_id = arguments.get("issue_id")
                    target_issue_id = arguments.get("target_issue_id")
                    if not all([issue_id, target_issue_id]):
                        raise ValueError("Missing required arguments: issue_id, target_issue_id")
                    result = bml_server.move_issue_above(issue_id, target_issue_id, arguments.get("repo"))
                
                case BMLTools.MOVE_ISSUE_BELOW.value:
                    issue_id = arguments.get("issue_id")
                    target_issue_id = arguments.get("target_issue_id")
                    if not all([issue_id, target_issue_id]):
                        raise ValueError("Missing required arguments: issue_id, target_issue_id")
                    result = bml_server.move_issue_below(issue_id, target_issue_id, arguments.get("repo"))
                
                case BMLTools.MOVE_ISSUE_BETWEEN.value:
                    issue_id = arguments.get("issue_id")
                    above_issue_id = arguments.get("above_issue_id")
                    below_issue_id = arguments.get("below_issue_id")
                    if not all([issue_id, above_issue_id, below_issue_id]):
                        raise ValueError("Missing required arguments: issue_id, above_issue_id, below_issue_id")
                    result = bml_server.move_issue_between(
                        issue_id, above_issue_id, below_issue_id, arguments.get("repo")
                    )
                
                case BMLTools.SET_ISSUE_STATUS.value:
                    issue_id = arguments.get("issue_id")
                    status = arguments.get("status")
                    if not all([issue_id, status]):
                        raise ValueError("Missing required arguments: issue_id, status")
                    result = bml_server.set_issue_status(issue_id, status, arguments.get("repo"))
                
                case BMLTools.VIEW_KANBAN.value:
                    result = bml_server.view_kanban(arguments.get("repo"))
                
                case BMLTools.GET_KANBAN_LANE.value:
                    status = arguments.get("status")
                    if not status:
                        raise ValueError("Missing required argument: status")
                    result = bml_server.get_kanban_lane(status, arguments.get("repo"))
                
                case _:
                    raise ValueError(f"Unknown tool: {name}")
            
            return [
                TextContent(type="text", text=json.dumps(result, indent=2))
            ]
        
        except Exception as e:
            raise ValueError(f"Error processing BML operation: {str(e)}")
    
    # Initialize server and run
    options = server.create_initialization_options()
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, options)