import os
import json
from typing import Optional, List, Dict, Any, Union

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

# Initialize Confluence client
CONFLUENCE_URL = os.getenv("CONFLUENCE_URL")
CONFLUENCE_USERNAME = os.getenv("CONFLUENCE_USERNAME")
CONFLUENCE_API_TOKEN = os.getenv("CONFLUENCE_API_TOKEN")

if not all([CONFLUENCE_URL, CONFLUENCE_USERNAME, CONFLUENCE_API_TOKEN]):
    raise ValueError("Missing required Confluence environment variables")

confluence = Confluence(
    url=CONFLUENCE_URL,
    username=CONFLUENCE_USERNAME,
    password=CONFLUENCE_API_TOKEN,
    cloud=True  # Assuming cloud instance, set to False for server
)

server = Server("hc-mcp-confluence")

@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    """List available Confluence tools."""
    return [
        types.Tool(
            name="get_page",
            description="Get content of a specific Confluence page",
            inputSchema={
                "type": "object",
                "properties": {
                    "page_id": {
                        "type": "string",
                        "description": "Confluence page ID"
                    },
                    "expand": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "List of properties to expand"
                    }
                },
                "required": ["page_id"]
            }
        ),
        types.Tool(
            name="create_page",
            description="Create a new Confluence page",
            inputSchema={
                "type": "object",
                "properties": {
                    "space_key": {
                        "type": "string",
                        "description": "Key of the space where the page will be created"
                    },
                    "title": {
                        "type": "string",
                        "description": "Title of the page"
                    },
                    "body": {
                        "type": "string",
                        "description": "Content of the page in storage format (XHTML)"
                    },
                    "parent_id": {
                        "type": "string",
                        "description": "ID of the parent page (optional)"
                    },
                    "labels": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "List of labels to add to the page"
                    }
                },
                "required": ["space_key", "title", "body"]
            }
        ),
        types.Tool(
            name="update_page",
            description="Update an existing Confluence page",
            inputSchema={
                "type": "object",
                "properties": {
                    "page_id": {
                        "type": "string",
                        "description": "ID of the page to update"
                    },
                    "title": {
                        "type": "string",
                        "description": "New title for the page"
                    },
                    "body": {
                        "type": "string",
                        "description": "New content for the page in storage format (XHTML)"
                    },
                    "version_comment": {
                        "type": "string",
                        "description": "Comment for the new version"
                    }
                },
                "required": ["page_id", "body"]
            }
        ),
        types.Tool(
            name="delete_page",
            description="Delete a Confluence page",
            inputSchema={
                "type": "object",
                "properties": {
                    "page_id": {
                        "type": "string",
                        "description": "ID of the page to delete"
                    }
                },
                "required": ["page_id"]
            }
        ),
        types.Tool(
            name="search_content",
            description="Search Confluence content using CQL",
            inputSchema={
                "type": "object",
                "properties": {
                    "cql": {
                        "type": "string",
                        "description": "Confluence Query Language (CQL) search string"
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Maximum number of results to return",
                        "default": 25
                    },
                    "expand": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "List of properties to expand in the results"
                    }
                },
                "required": ["cql"]
            }
        ),
        types.Tool(
            name="get_space",
            description="Get details about a Confluence space",
            inputSchema={
                "type": "object",
                "properties": {
                    "space_key": {
                        "type": "string",
                        "description": "Key of the space"
                    },
                    "expand": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "List of properties to expand"
                    }
                },
                "required": ["space_key"]
            }
        ),
        types.Tool(
            name="list_spaces",
            description="List all Confluence spaces",
            inputSchema={
                "type": "object",
                "properties": {
                    "limit": {
                        "type": "integer",
                        "description": "Maximum number of spaces to return",
                        "default": 25
                    },
                    "expand": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "List of properties to expand"
                    }
                }
            }
        ),
        types.Tool(
            name="add_page_labels",
            description="Add labels to a Confluence page",
            inputSchema={
                "type": "object",
                "properties": {
                    "page_id": {
                        "type": "string",
                        "description": "ID of the page"
                    },
                    "labels": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "List of labels to add"
                    }
                },
                "required": ["page_id", "labels"]
            }
        ),
        types.Tool(
            name="get_page_labels",
            description="Get labels of a Confluence page",
            inputSchema={
                "type": "object",
                "properties": {
                    "page_id": {
                        "type": "string",
                        "description": "ID of the page"
                    }
                },
                "required": ["page_id"]
            }
        ),
        types.Tool(
            name="get_page_attachments",
            description="Get attachments of a Confluence page",
            inputSchema={
                "type": "object",
                "properties": {
                    "page_id": {
                        "type": "string",
                        "description": "ID of the page"
                    }
                },
                "required": ["page_id"]
            }
        )
    ]

def format_response(result: Dict[str, Any]) -> str:
    """Format API response with proper structure and indentation."""
    return json.dumps(result, indent=2)

def enhance_page_response(page: Dict[str, Any], expand: List[str]) -> Dict[str, Any]:
    """Enhance page response with proper structure and expanded content."""
    if not page:
        return page

    # Add body.storage if requested
    if expand and "body.storage" in expand and "body" not in page:
        page["body"] = {
            "storage": {
                "value": page.get("content", ""),
                "representation": "storage",
                "embeddedContent": [],
                "_expandable": {
                    "content": f"/rest/api/content/{page.get('id', '')}"
                }
            },
            "_expandable": {
                "editor": "",
                "atlas_doc_format": "",
                "view": "",
                "export_view": "",
                "styled_view": "",
                "dynamic": "",
                "editor2": "",
                "anonymous_export_view": ""
            }
        }

    return page

def enhance_space_response(space: Dict[str, Any]) -> Dict[str, Any]:
    """Enhance space response with proper structure."""
    if not space:
        return space

    # Add description if not present
    if "description" not in space:
        space["description"] = {
            "_expandable": {
                "view": "",
                "plain": ""
            }
        }

    return space

@server.call_tool()
async def handle_call_tool(
    name: str, arguments: dict | None
) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
    """Handle Confluence tool execution requests."""
    if not arguments:
        raise ValueError("Missing arguments")

    try:
        if name == "get_page":
            page_id = arguments["page_id"]
            expand = arguments.get("expand", [])
            
            # First get the basic page info
            result = confluence.get_page_by_id(page_id)
            
            # If body.storage is requested, get it in a separate call
            if expand and "body.storage" in expand:
                body_content = confluence.get_page_by_id(
                    page_id,
                    expand="body.storage"
                )
                if "body" in body_content:
                    result["body"] = body_content["body"]
                    if "storage" in result["body"]:
                        result["body"]["storage"]["embeddedContent"] = []
                        result["body"]["storage"]["_expandable"] = {
                            "content": f"/rest/api/content/{page_id}"
                        }
                    result["body"]["_expandable"] = {
                        "editor": "",
                        "atlas_doc_format": "",
                        "view": "",
                        "export_view": "",
                        "styled_view": "",
                        "dynamic": "",
                        "editor2": "",
                        "anonymous_export_view": ""
                    }
            
            return [types.TextContent(type="text", text=format_response(result))]

        elif name == "create_page":
            space_key = arguments["space_key"]
            title = arguments["title"]
            body = arguments["body"]
            parent_id = arguments.get("parent_id")
            labels = arguments.get("labels", [])
            
            result = confluence.create_page(
                space=space_key,
                title=title,
                body=body,
                parent_id=parent_id,
                representation="storage"
            )
            
            if labels:
                confluence.set_page_labels(result["id"], labels)
            
            # Get complete page details after creation
            complete_result = confluence.get_page_by_id(
                result["id"],
                expand="body.storage,version"
            )
            return [types.TextContent(type="text", text=format_response(complete_result))]

        elif name == "update_page":
            page_id = arguments["page_id"]
            body = arguments["body"]
            title = arguments.get("title")
            version_comment = arguments.get("version_comment")
            
            page = confluence.get_page_by_id(page_id)
            if title is None:
                title = page["title"]
            
            result = confluence.update_page(
                page_id=page_id,
                title=title,
                body=body,
                version_comment=version_comment
            )
            
            # Get complete page details after update
            complete_result = confluence.get_page_by_id(
                page_id,
                expand="body.storage,version"
            )
            return [types.TextContent(type="text", text=format_response(complete_result))]

        elif name == "delete_page":
            page_id = arguments["page_id"]
            result = confluence.remove_page(page_id)
            return [types.TextContent(type="text", text=format_response({"status": "success", "message": "Page deleted"}))]

        elif name == "search_content":
            cql = arguments["cql"]
            limit = arguments.get("limit", 25)
            expand = arguments.get("expand", [])
            
            # Always include body.storage in expand
            if "body.storage" not in expand:
                expand.append("body.storage")
            
            result = confluence.cql(
                cql,
                limit=limit,
                expand=",".join(expand)
            )
            
            # Enhance search results
            if "results" in result:
                for item in result["results"]:
                    if "content" in item:
                        item["content"] = enhance_page_response(item["content"], expand)
            
            return [types.TextContent(type="text", text=format_response(result))]

        elif name == "get_space":
            space_key = arguments["space_key"]
            expand = arguments.get("expand", [])
            
            result = confluence.get_space(
                space_key,
                expand=",".join(expand) if expand else None
            )
            
            enhanced_result = enhance_space_response(result)
            return [types.TextContent(type="text", text=format_response(enhanced_result))]

        elif name == "list_spaces":
            limit = arguments.get("limit", 25)
            expand = arguments.get("expand", [])
            
            # Always include description in expand
            if "description" not in expand:
                expand.append("description")
            
            result = confluence.get_all_spaces(
                limit=limit,
                expand=",".join(expand)
            )
            
            # Enhance space results
            if "results" in result:
                result["results"] = [
                    enhance_space_response(space)
                    for space in result["results"]
                ]
            
            return [types.TextContent(type="text", text=format_response(result))]

        elif name == "add_page_labels":
            page_id = arguments["page_id"]
            labels = arguments["labels"]
            result = confluence.set_page_labels(page_id, labels)
            return [types.TextContent(type="text", text=format_response(result))]

        elif name == "get_page_labels":
            page_id = arguments["page_id"]
            result = confluence.get_page_labels(page_id)
            return [types.TextContent(type="text", text=format_response(result))]

        elif name == "get_page_attachments":
            page_id = arguments["page_id"]
            result = confluence.get_attachments_from_content(page_id)
            return [types.TextContent(type="text", text=format_response(result))]

        else:
            raise ValueError(f"Unknown tool: {name}")

    except Exception as e:
        error_response = {
            "statusCode": 500,
            "error": str(e),
            "details": {
                "name": e.__class__.__name__,
                "message": str(e),
                "tool": name,
                "arguments": arguments
            }
        }
        return [types.TextContent(type="text", text=format_response(error_response))]

async def main():
    """Run the Confluence MCP server."""
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="hc-mcp-confluence",
                server_version="0.1.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )
