import importlib
import json
import logging
import os
from typing import Any, Callable, Dict, List, Optional, Union

import uvicorn
from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel, Field

from .models import EvaluateResult

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class Message(BaseModel):
    """Model for a conversation message."""

    role: str
    content: str

    class Config:
        extra = "allow"  # Allow extra fields


class RewardRequest(BaseModel):
    """Request model for reward endpoints."""

    messages: List[Message] = Field(..., description="List of conversation messages")
    ground_truth: Optional[Union[str, List[Message]]] = Field(
        None, description="Ground truth data (string or list of messages) for context"
    )

    class Config:
        extra = "allow"  # Allow extra fields for arbitrary kwargs


class RewardServer:
    """
    Server for hosting reward functions.

    This class creates a FastAPI server that can host reward functions.

    Args:
        func_path: Path to the reward function to host (e.g., "module.path:function_name")
        host: Host to bind the server to
        port: Port to bind the server to
    """

    def __init__(
        self,
        func_path: str,
        host: str = "0.0.0.0",
        port: int = 8000,
    ):
        self.func_path = func_path
        self.host = host
        self.port = port
        self.app = FastAPI(title="Reward Function Server")

        # Load the reward function
        self.reward_func = self._load_function()

        # Register the endpoints
        self._setup_routes()

    def _load_function(self):
        """Load the reward function from the provided path."""
        try:
            if ":" not in self.func_path:
                raise ValueError(
                    f"Invalid func_path format: {self.func_path}, expected 'module.path:function_name'"
                )

            module_path, func_name = self.func_path.split(":", 1)
            module = importlib.import_module(module_path)
            func = getattr(module, func_name)

            logger.info(f"Loaded reward function {func_name} from {module_path}")
            return func
        except (ImportError, AttributeError) as e:
            raise ImportError(
                f"Failed to load function from path {self.func_path}: {str(e)}"
            )

    def _setup_routes(self):
        """Set up the API routes."""

        @self.app.get("/")
        async def root():
            """Get server info."""
            return {
                "status": "ok",
                "reward_function": self.func_path,
                "endpoints": ["/reward"],
            }

        @self.app.post("/reward")
        async def reward(request: RewardRequest):
            """
            Get reward score for messages.

            Args:
                request: RewardRequest object with messages and optional parameters

            Returns:
                EvaluateResult object with score and metrics
            """
            try:
                # Extract kwargs from the request
                kwargs = request.dict(exclude={"messages", "ground_truth"})

                # Set default for ground_truth if not provided and expected as list
                ground_truth_data = request.ground_truth
                if ground_truth_data is None:
                    # This default applies if ground_truth is expected to be a list of messages for context
                    ground_truth_data = (
                        request.messages[:-1] if request.messages else []
                    )

                # Call the reward function
                result = self.reward_func(
                    messages=request.messages,
                    ground_truth=ground_truth_data,
                    **kwargs,
                )

                # Handle different return types
                # The self.reward_func is expected to be decorated by the new @reward_function,
                # which returns a dictionary.
                if isinstance(result, dict) and "score" in result:
                    return result
                elif isinstance(
                    result, EvaluateResult
                ):  # Should not happen if func is from new decorator
                    logger.warning(
                        "Reward function returned EvaluateResult object directly to server; expected dict."
                    )
                    return result.model_dump()
                elif isinstance(result, tuple) and len(result) == 2:  # Legacy tuple
                    logger.warning(
                        "Reward function returned legacy tuple format to server."
                    )
                    score, components = result
                    return {"score": score, "metrics": components}
                else:
                    raise TypeError(
                        f"Invalid return type from reward function after decoration: {type(result)}"
                    )

            except Exception as e:
                logger.error(f"Error processing reward request: {str(e)}")
                raise HTTPException(status_code=500, detail=str(e))

        @self.app.get("/health")
        async def health():
            """Health check endpoint."""
            return {"status": "ok"}

    def run(self):
        """Run the server."""
        logger.info(f"Starting reward server on {self.host}:{self.port}")
        uvicorn.run(self.app, host=self.host, port=self.port)


def serve(func_path: str, host: str = "0.0.0.0", port: int = 8000):
    """
    Serve a reward function as an HTTP API.

    Args:
        func_path: Path to the reward function to serve (e.g., "module.path:function_name")
        host: Host to bind the server to
        port: Port to bind the server to
    """
    server = RewardServer(func_path=func_path, host=host, port=port)
    server.run()


# ngrok-based serve_tunnel is deprecated in favor of Serveo via subprocess_manager.
# def serve_tunnel(func_path: str, port: int = 8000):
#     """
#     Serve a reward function with an ngrok tunnel.
#     DEPRECATED.
#     """
#     try:
#         import pyngrok.ngrok as ngrok  # type: ignore
#     except ImportError:
#         raise ImportError(
#             "The 'pyngrok' package is required to use serve_tunnel. "
#             "Please install it with 'pip install pyngrok'."
#         )
#
#     # Open the tunnel
#     tunnel = ngrok.connect(port)
#     public_url = tunnel.public_url
#
#     # Print the tunnel URL
#     logger.info(f"Reward function available at: {public_url}/reward")
#
#     # Start the server
#     serve(func_path=func_path, host="0.0.0.0", port=port)


def create_app(reward_func: Callable[..., EvaluateResult]) -> FastAPI:
    """
    Create a FastAPI app for the given reward function.

    This function creates a FastAPI app that can be used to serve a reward function.
    It's particularly useful for testing or when you want to manage the lifecycle
    of the app yourself.

    Args:
        reward_func: The reward function to serve

    Returns:
        A FastAPI app instance
    """
    app = FastAPI(title="Reward Function Server")

    @app.get("/")
    async def root():
        """Get server info."""
        return {"status": "ok", "endpoints": ["/reward"]}

    @app.post("/reward")
    async def reward(request_data: RewardRequest):
        """
        Get reward score for messages.

        Args:
            request_data: RewardRequest object with messages and optional parameters

        Returns:
            EvaluateResult object with score and metrics
        """
        try:
            # Convert Pydantic models to dictionaries using model_dump (Pydantic v2)
            messages = [msg.model_dump() for msg in request_data.messages]
            ground_truth_data: Optional[Union[str, List[Dict[str, Any]]]] = None

            if isinstance(request_data.ground_truth, str):
                ground_truth_data = request_data.ground_truth
            elif isinstance(request_data.ground_truth, list):
                ground_truth_data = [
                    msg.model_dump() for msg in request_data.ground_truth
                ]

            # Extract kwargs from any extra fields
            kwargs = {
                k: v
                for k, v in request_data.model_dump().items()
                if k not in ["messages", "ground_truth"]
            }

            # Set default for ground_truth if not provided and expected as list
            if ground_truth_data is None:
                # This default applies if ground_truth is expected to be a list of messages for context
                ground_truth_data = messages[:-1] if messages else []

            # Call the reward function
            result = reward_func(
                messages=messages, ground_truth=ground_truth_data, **kwargs
            )

            # Handle different return types
            # The reward_func is expected to be decorated by the new @reward_function,
            # which returns a dictionary.
            if isinstance(result, dict) and "score" in result:
                return result
            elif isinstance(
                result, EvaluateResult
            ):  # Should not happen if func is from new decorator
                logger.warning(
                    "Reward function passed to create_app returned EvaluateResult object directly; expected dict after decoration."
                )
                return result.model_dump()
            elif isinstance(result, tuple) and len(result) == 2:  # Legacy tuple
                logger.warning(
                    "Reward function passed to create_app returned legacy tuple format."
                )
                score, components = result
                return {"score": score, "metrics": components}
            else:
                raise TypeError(
                    f"Invalid return type from reward function after decoration: {type(result)}"
                )

        except Exception as e:
            logger.error(f"Error processing reward request: {str(e)}")
            raise HTTPException(status_code=500, detail=str(e))

    @app.get("/health")
    async def health():
        """Health check endpoint."""
        return {"status": "ok"}

    return app
