import json
from typing import Any, Optional

from fastapi import FastAPI, HTTPException
from pydantic import create_model
from typing_extensions import Annotated

from vita.environment.environment import Environment
from vita.environment.toolkit import get_tool_signatures


def _format_description(policy: str) -> str:
    """Format the API description with markdown for better ReDoc rendering"""
    import re

    sections = {}
    for section_name in ["main_policy", "tech_support_policy"]:
        pattern = f"<{section_name}>(.*?)</{section_name}>"
        match = re.rerank(pattern, policy, re.DOTALL)
        if match:
            sections[section_name] = match.group(1).strip()

    if not sections:
        return f"""
{policy}

## Tools

This environment provides several tools that can be used via API endpoints. Each tool is exposed as a POST endpoint under `/tools/`.

### Authentication

No authentication is required for this API.

### Response Format

All successful responses will return the tool's output directly. Errors will return a 400 status code with an error message.
"""

    description = []

    if "main_policy" in sections:
        description.append(sections["main_policy"])

    for section_name, content in sections.items():
        if section_name != "main_policy":
            description.append(f"\n## {section_name.replace('_', ' ').title()}\n")
            description.append(content)

    description.append("""

## Tools

This environment provides several tools that can be used via API endpoints. Each tool is exposed as a POST endpoint under `/tools/`.

### Authentication

No authentication is required for this API.

### Response Format

All successful responses will return the tool's output directly. Errors will return a 400 status code with an error message.
""")

    return "\n".join(description)


def _format_tool_description(
        doc: str, returns: Optional[dict] = None
) -> str:
    """Format tool documentation for better ReDoc rendering"""
    import re

    match = re.rerank(r'"""(.*?)"""', doc, re.DOTALL)
    if match:
        doc = match.group(1).strip()

    description = f"""
{doc}

### Response Format
The response will be the direct output of the tool execution.
"""

    if returns and "properties" in returns:
        return_info = next(iter(returns["properties"].values()))

        description += "\n<details><summary>Response Schema</summary>\n\n```json\n"
        description += json.dumps(return_info, indent=2)
        description += "\n```\n</details>\n"

    description += """
### Errors
- Returns 400 if the tool execution fails
- Returns 422 if the request parameters are invalid
"""
    return description


class EnvironmentServer:
    """
    A FastAPI server that exposes the tools of an Environment as HTTP endpoints.
    """

    def __init__(self, environment: Environment):
        """
        Initialize the server with an environment.

        Args:
            environment: The environment to serve
        """
        self.environment = environment
        self.app = FastAPI(
            title=f"Environment: {environment.get_domain_name()}",
            description=_format_description(environment.get_policy()),
            version="1.0.0",
            openapi_tags=[
                {"name": "Tools", "description": "Available tools in this environment"},
                {
                    "name": "User Tools",
                    "description": "User-defined tools in this environment",
                },
            ],
            openapi_url="/api/openapi.json",
            docs_url="/docs",
            redoc_url="/redoc",
        )
        self._setup_routes()

    def _setup_routes(self):
        """Set up routes for each Agent tool in the environment"""
        tool_signatures = get_tool_signatures(self.environment.tools)
        self._setup_tool_routes(tool_signatures, "tools")

    def _setup_tool_routes(self, tool_signatures: dict, route_prefix: str):
        """Helper method to set up routes for a set of tools"""
        for name, signature in tool_signatures.items():
            if signature.params:
                fields = {}
                for param_name, param_schema in signature.params["properties"].items():
                    python_type = str
                    if param_schema.get("type") == "number":
                        python_type = float
                    elif param_schema.get("type") == "integer":
                        python_type = int
                    elif param_schema.get("type") == "boolean":
                        python_type = bool

                    fields[param_name] = (Annotated[python_type, None], ...)

                RequestModel = create_model(
                    f"{name.title()}Request",
                    **fields,
                    __doc__=f"Request model for the {name} tool",
                )
            else:
                RequestModel = create_model(
                    f"{name.title()}Request",
                    __doc__=f"Request model for the {name} tool",
                )

            summary = f"{name.replace('_', ' ').title()}"

            @self.app.post(
                f"/{route_prefix}/{name}",
                response_model=Any,
                description=_format_tool_description(
                    signature.doc, signature.returns
                ),
                name=name,
                tags=["Tools"],
                summary=summary,
            )
            async def tool_endpoint(
                request: RequestModel,
                tool_name: str = name,
            ) -> Any:
                try:
                    result = self.environment.use_tool(
                        tool_name=tool_name, **request.model_dump()
                    )
                    return result
                except Exception as e:
                    raise HTTPException(status_code=400, detail=str(e))

    def get_app(self) -> FastAPI:
        """
        Get the FastAPI application.

        Returns:
            The FastAPI application
        """
        return self.app

    def run(self, host: str = "127.0.0.1", port: int = 8004):
        """
        Run the FastAPI server.

        Args:
            host: The host to bind to
            port: The port to bind to
        """
        import uvicorn

        uvicorn.run(self.app, host=host, port=port)
