#!/usr/bin/python
import sys

# coding: utf-8
import json
import os
import argparse
import logging
import uvicorn
from typing import Optional, Any, List
from contextlib import asynccontextmanager

from pydantic_ai import Agent, ModelSettings, RunContext
from arr_mcp.utils import (
    to_integer,
    to_boolean,
    to_float,
    to_list,
    to_dict,
    get_mcp_config_path,
    get_skills_path,
    load_skills_from_directory,
    create_model,
    prune_large_messages,
)

from fastapi import FastAPI, Request
from starlette.responses import Response, StreamingResponse
from pydantic import ValidationError
from pydantic_ai.ui import SSE_CONTENT_TYPE
from pydantic_ai.ui.ag_ui import AGUIAdapter

from arr_mcp.lidarr_agent import create_agent as create_lidarr_agent
from arr_mcp.sonarr_agent import create_agent as create_sonarr_agent
from arr_mcp.radarr_agent import create_agent as create_radarr_agent
from arr_mcp.prowlarr_agent import create_agent as create_prowlarr_agent
from arr_mcp.chaptarr_agent import create_agent as create_chaptarr_agent
from arr_mcp.seerr_agent import create_agent as create_seerr_agent
from arr_mcp.bazarr_agent import create_agent as create_bazarr_agent

__version__ = "0.2.7"

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler()],
)
logging.getLogger("pydantic_ai").setLevel(logging.INFO)
logging.getLogger("fastmcp").setLevel(logging.INFO)
logging.getLogger("httpx").setLevel(logging.INFO)
logger = logging.getLogger(__name__)

DEFAULT_HOST = os.getenv("HOST", "0.0.0.0")
DEFAULT_PORT = to_integer(os.getenv("PORT", "9000"))
DEFAULT_DEBUG = to_boolean(os.getenv("DEBUG", "False"))
DEFAULT_PROVIDER = os.getenv("PROVIDER", "openai")
DEFAULT_MODEL_ID = os.getenv("MODEL_ID", "qwen/qwen3-coder-next")
DEFAULT_LLM_BASE_URL = os.getenv("LLM_BASE_URL", "http://host.docker.internal:1234/v1")
DEFAULT_LLM_API_KEY = os.getenv("LLM_API_KEY", "ollama")
DEFAULT_MCP_URL = os.getenv("MCP_URL", None)
DEFAULT_MCP_CONFIG = os.getenv("MCP_CONFIG", get_mcp_config_path())
DEFAULT_SKILLS_DIRECTORY = os.getenv("SKILLS_DIRECTORY", get_skills_path())
DEFAULT_ENABLE_WEB_UI = to_boolean(os.getenv("ENABLE_WEB_UI", "False"))
DEFAULT_SSL_VERIFY = to_boolean(os.getenv("SSL_VERIFY", "True"))

DEFAULT_MAX_TOKENS = to_integer(os.getenv("MAX_TOKENS", "16384"))
DEFAULT_TEMPERATURE = to_float(os.getenv("TEMPERATURE", "0.7"))
DEFAULT_TOP_P = to_float(os.getenv("TOP_P", "1.0"))
DEFAULT_TIMEOUT = to_float(os.getenv("TIMEOUT", "32400.0"))
DEFAULT_TOOL_TIMEOUT = to_float(os.getenv("TOOL_TIMEOUT", "32400.0"))
DEFAULT_PARALLEL_TOOL_CALLS = to_boolean(os.getenv("PARALLEL_TOOL_CALLS", "True"))
DEFAULT_SEED = to_integer(os.getenv("SEED", None))
DEFAULT_PRESENCE_PENALTY = to_float(os.getenv("PRESENCE_PENALTY", "0.0"))
DEFAULT_FREQUENCY_PENALTY = to_float(os.getenv("FREQUENCY_PENALTY", "0.0"))
DEFAULT_LOGIT_BIAS = to_dict(os.getenv("LOGIT_BIAS", None))
DEFAULT_STOP_SEQUENCES = to_list(os.getenv("STOP_SEQUENCES", None))
DEFAULT_EXTRA_HEADERS = to_dict(os.getenv("EXTRA_HEADERS", None))
DEFAULT_EXTRA_BODY = to_dict(os.getenv("EXTRA_BODY", None))

AGENT_NAME = "ArrAgent"
AGENT_DESCRIPTION = (
    "A supervisor agent for the Arr stack (Lidarr, Sonarr, Radarr, Prowlarr, Chaptarr)."
)


SUPERVISOR_SYSTEM_PROMPT = os.environ.get(
    "SUPERVISOR_SYSTEM_PROMPT",
    default=(
        "You are the Arr Supervisor Agent.\n"
        "Your goal is to assist the user by assigning tasks to specialized child agents for each service in the Arr stack.\n"
        "You have access to the following specialists:\n"
        "- Lidarr Agent: For music management.\n"
        "- Sonarr Agent: For TV show management.\n"
        "- Radarr Agent: For movie management.\n"
        "- Prowlarr Agent: For indexer management.\n"
        "- Chaptarr Agent: For book management.\n"
        "- Seerr Agent: For media requests and discovery.\n"
        "- Bazarr Agent: For subtitle management.\n"
        "\n"
        "Analyze the user's request and determine which service it relates to.\n"
        "Then, use the appropriate 'ask_..._agent' tool to delegate the task.\n"
        "Synthesize the results from the child agent into a final helpful response.\n"
        "Always be warm, professional, and helpful."
    ),
)


def create_agent(
    provider: str = DEFAULT_PROVIDER,
    model_id: str = DEFAULT_MODEL_ID,
    base_url: Optional[str] = DEFAULT_LLM_BASE_URL,
    api_key: Optional[str] = DEFAULT_LLM_API_KEY,
    mcp_url: str = DEFAULT_MCP_URL,
    mcp_config: str = DEFAULT_MCP_CONFIG,
    skills_directory: Optional[str] = DEFAULT_SKILLS_DIRECTORY,
    ssl_verify: bool = DEFAULT_SSL_VERIFY,
) -> Agent:
    """
    Creates the Supervisor Agent with sub-agents registered as tools.
    """
    logger.info("Initializing Arr Supervisor Agent...")

    model = create_model(
        provider=provider,
        model_id=model_id,
        base_url=base_url,
        api_key=api_key,
        ssl_verify=ssl_verify,
        timeout=DEFAULT_TIMEOUT,
    )
    settings = ModelSettings(
        max_tokens=DEFAULT_MAX_TOKENS,
        temperature=DEFAULT_TEMPERATURE,
        top_p=DEFAULT_TOP_P,
        timeout=DEFAULT_TIMEOUT,
        parallel_tool_calls=DEFAULT_PARALLEL_TOOL_CALLS,
        seed=DEFAULT_SEED,
        presence_penalty=DEFAULT_PRESENCE_PENALTY,
        frequency_penalty=DEFAULT_FREQUENCY_PENALTY,
        logit_bias=DEFAULT_LOGIT_BIAS,
        stop_sequences=DEFAULT_STOP_SEQUENCES,
        extra_headers=DEFAULT_EXTRA_HEADERS,
        extra_body=DEFAULT_EXTRA_BODY,
    )

    child_agents = {}

    try:
        child_agents["lidarr"] = create_lidarr_agent(
            provider,
            model_id,
            base_url,
            api_key,
            mcp_url,
            mcp_config,
            skills_directory,
            ssl_verify=ssl_verify,
            timeout=DEFAULT_TIMEOUT,
        )
        logger.info("Lidarr Agent initialized.")
    except Exception as e:
        logger.error(f"Failed to initialize Lidarr Agent: {e}")

    try:
        child_agents["sonarr"] = create_sonarr_agent(
            provider,
            model_id,
            base_url,
            api_key,
            mcp_url,
            mcp_config,
            skills_directory,
            ssl_verify=ssl_verify,
            timeout=DEFAULT_TIMEOUT,
        )
        logger.info("Sonarr Agent initialized.")
    except Exception as e:
        logger.error(f"Failed to initialize Sonarr Agent: {e}")

    try:
        child_agents["radarr"] = create_radarr_agent(
            provider,
            model_id,
            base_url,
            api_key,
            mcp_url,
            mcp_config,
            skills_directory,
            ssl_verify=ssl_verify,
            timeout=DEFAULT_TIMEOUT,
        )
        logger.info("Radarr Agent initialized.")
    except Exception as e:
        logger.error(f"Failed to initialize Radarr Agent: {e}")

    try:
        child_agents["prowlarr"] = create_prowlarr_agent(
            provider,
            model_id,
            base_url,
            api_key,
            mcp_url,
            mcp_config,
            skills_directory,
            ssl_verify=ssl_verify,
            timeout=DEFAULT_TIMEOUT,
        )
        logger.info("Prowlarr Agent initialized.")
    except Exception as e:
        logger.error(f"Failed to initialize Prowlarr Agent: {e}")

    try:
        child_agents["chaptarr"] = create_chaptarr_agent(
            provider,
            model_id,
            base_url,
            api_key,
            mcp_url,
            mcp_config,
            skills_directory,
            ssl_verify=ssl_verify,
            timeout=DEFAULT_TIMEOUT,
        )
        logger.info("Chaptarr Agent initialized.")
    except Exception as e:
        logger.error(f"Failed to initialize Chaptarr Agent: {e}")

    try:
        child_agents["seerr"] = create_seerr_agent(
            provider,
            model_id,
            base_url,
            api_key,
            mcp_url,
            mcp_config,
            skills_directory,
            ssl_verify=ssl_verify,
            timeout=DEFAULT_TIMEOUT,
        )
        logger.info("Seerr Agent initialized.")
    except Exception as e:
        logger.error(f"Failed to initialize Seerr Agent: {e}")

    try:
        child_agents["bazarr"] = create_bazarr_agent(
            provider,
            model_id,
            base_url,
            api_key,
            mcp_url,
            mcp_config,
            skills_directory,
            ssl_verify=ssl_verify,
            timeout=DEFAULT_TIMEOUT,
        )
        logger.info("Bazarr Agent initialized.")
    except Exception as e:
        logger.error(f"Failed to initialize Bazarr Agent: {e}")

    supervisor = Agent(
        name=AGENT_NAME,
        system_prompt=SUPERVISOR_SYSTEM_PROMPT,
        model=model,
        model_settings=settings,
        deps_type=Any,
    )

    @supervisor.tool
    async def ask_lidarr_agent(ctx: RunContext[Any], task: str) -> str:
        """
        Delegate a task to the Lidarr Agent (Music).
        Use this for anything related to artists, albums, tracks, or music management.
        """
        if "lidarr" not in child_agents:
            return "Lidarr Agent is not available."
        return (
            await child_agents["lidarr"].run(task, usage=ctx.usage, deps=ctx.deps)
        ).output

    @supervisor.tool
    async def ask_sonarr_agent(ctx: RunContext[Any], task: str) -> str:
        """
        Delegate a task to the Sonarr Agent (TV Shows).
        Use this for anything related to series, episodes, seasons, or TV management.
        """
        if "sonarr" not in child_agents:
            return "Sonarr Agent is not available."
        return (
            await child_agents["sonarr"].run(task, usage=ctx.usage, deps=ctx.deps)
        ).output

    @supervisor.tool
    async def ask_radarr_agent(ctx: RunContext[Any], task: str) -> str:
        """
        Delegate a task to the Radarr Agent (Movies).
        Use this for anything related to movies, films, collections, or cinema management.
        """
        if "radarr" not in child_agents:
            return "Radarr Agent is not available."
        return (
            await child_agents["radarr"].run(task, usage=ctx.usage, deps=ctx.deps)
        ).output

    @supervisor.tool
    async def ask_prowlarr_agent(ctx: RunContext[Any], task: str) -> str:
        """
        Delegate a task to the Prowlarr Agent (Indexers).
        Use this for managing indexers, trackers, or download clients configuration.
        """
        if "prowlarr" not in child_agents:
            return "Prowlarr Agent is not available."
        return (
            await child_agents["prowlarr"].run(task, usage=ctx.usage, deps=ctx.deps)
        ).output

    @supervisor.tool
    async def ask_chaptarr_agent(ctx: RunContext[Any], task: str) -> str:
        """
        Delegate a task to the Chaptarr Agent (Books).
        Use this for anything related to books, authors, or ebook management.
        """
        if "chaptarr" not in child_agents:
            return "Chaptarr Agent is not available."
        return (
            await child_agents["chaptarr"].run(task, usage=ctx.usage, deps=ctx.deps)
        ).output

    @supervisor.tool
    async def ask_seerr_agent(ctx: RunContext[Any], task: str) -> str:
        """
        Delegate a task to the Seerr Agent (Requests & Discovery).
        Use this for searching movies/TV, making requests, or managing users.
        """
        if "seerr" not in child_agents:
            return "Seerr Agent is not available."
        return (
            await child_agents["seerr"].run(task, usage=ctx.usage, deps=ctx.deps)
        ).output

    @supervisor.tool
    async def ask_bazarr_agent(ctx: RunContext[Any], task: str) -> str:
        """
        Delegate a task to the Bazarr Agent (Subtitles).
        Use this for finding, downloading, or managing subtitles for movies and TV shows.
        """
        if "bazarr" not in child_agents:
            return "Bazarr Agent is not available."
        return (
            await child_agents["bazarr"].run(task, usage=ctx.usage, deps=ctx.deps)
        ).output

    return supervisor


async def chat(agent: Agent, prompt: str):
    result = await agent.run(prompt)
    print(f"Response:\n\n{result.output}")


async def node_chat(agent: Agent, prompt: str) -> List:
    nodes = []
    async with agent.iter(prompt) as agent_run:
        async for node in agent_run:
            nodes.append(node)
            print(node)
    return nodes


async def stream_chat(agent: Agent, prompt: str) -> None:
    async with agent.run_stream(prompt) as result:
        async for text_chunk in result.stream_text(delta=True):
            print(text_chunk, end="", flush=True)
        print("\nDone!")


def create_agent_server(
    provider: str = DEFAULT_PROVIDER,
    model_id: str = DEFAULT_MODEL_ID,
    base_url: Optional[str] = DEFAULT_LLM_BASE_URL,
    api_key: Optional[str] = DEFAULT_LLM_API_KEY,
    mcp_url: str = DEFAULT_MCP_URL,
    mcp_config: str = DEFAULT_MCP_CONFIG,
    skills_directory: Optional[str] = DEFAULT_SKILLS_DIRECTORY,
    debug: Optional[bool] = DEFAULT_DEBUG,
    host: Optional[str] = DEFAULT_HOST,
    port: Optional[int] = DEFAULT_PORT,
    enable_web_ui: bool = DEFAULT_ENABLE_WEB_UI,
    ssl_verify: bool = DEFAULT_SSL_VERIFY,
):
    print(
        f"Starting {AGENT_NAME}:"
        f"\tprovider={provider}"
        f"\tmodel={model_id}"
        f"\tbase_url={base_url}"
        f"\tmcp={mcp_url} | {mcp_config}"
        f"\tssl_verify={ssl_verify}"
    )
    agent = create_agent(
        provider=provider,
        model_id=model_id,
        base_url=base_url,
        api_key=api_key,
        mcp_url=mcp_url,
        mcp_config=mcp_config,
        skills_directory=skills_directory,
        ssl_verify=ssl_verify,
        timeout=DEFAULT_TIMEOUT,
    )

    skills = []
    if skills_directory and os.path.exists(skills_directory):
        skills = load_skills_from_directory(skills_directory)
        logger.info(f"Loaded {len(skills)} skills from {skills_directory}")

    a2a_app = agent.to_a2a(
        name=AGENT_NAME,
        description=AGENT_DESCRIPTION,
        version=__version__,
        skills=skills,
        debug=debug,
    )

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        if hasattr(a2a_app, "router") and hasattr(a2a_app.router, "lifespan_context"):
            async with a2a_app.router.lifespan_context(a2a_app):
                yield
        else:
            yield

    app = FastAPI(
        title=f"{AGENT_NAME} - A2A + AG-UI Server",
        description=AGENT_DESCRIPTION,
        debug=debug,
        lifespan=lifespan,
    )

    @app.get("/health")
    async def health_check():
        return {"status": "OK"}

    app.mount("/a2a", a2a_app)

    @app.post("/ag-ui")
    async def ag_ui_endpoint(request: Request) -> Response:
        accept = request.headers.get("accept", SSE_CONTENT_TYPE)
        try:
            run_input = AGUIAdapter.build_run_input(await request.body())
        except ValidationError as e:
            return Response(
                content=json.dumps(e.json()),
                media_type="application/json",
                status_code=422,
            )

        if hasattr(run_input, "messages"):
            run_input.messages = prune_large_messages(run_input.messages)

        adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
        event_stream = adapter.run_stream()
        sse_stream = adapter.encode_stream(event_stream)

        return StreamingResponse(
            sse_stream,
            media_type=accept,
        )

    if enable_web_ui:
        web_ui = agent.to_web(instructions=SUPERVISOR_SYSTEM_PROMPT)
        app.mount("/", web_ui)
        logger.info(
            "Starting server on %s:%s (A2A at /a2a, AG-UI at /ag-ui, Web UI: %s)",
            host,
            port,
            "Enabled at /" if enable_web_ui else "Disabled",
        )

    uvicorn.run(
        app,
        host=host,
        port=port,
        timeout_keep_alive=1800,
        timeout_graceful_shutdown=60,
        log_level="debug" if debug else "info",
    )


def agent_server():
    print(f"audio_transcriber_agent v{__version__}")
    parser = argparse.ArgumentParser(
        add_help=False, description=f"Run the {AGENT_NAME} A2A + AG-UI Server"
    )
    parser.add_argument(
        "--host", default=DEFAULT_HOST, help="Host to bind the server to"
    )
    parser.add_argument(
        "--port", type=int, default=DEFAULT_PORT, help="Port to bind the server to"
    )
    parser.add_argument("--debug", type=bool, default=DEFAULT_DEBUG, help="Debug mode")
    parser.add_argument("--reload", action="store_true", help="Enable auto-reload")

    parser.add_argument(
        "--provider",
        default=DEFAULT_PROVIDER,
        choices=["openai", "anthropic", "google", "huggingface"],
        help="LLM Provider",
    )
    parser.add_argument("--model-id", default=DEFAULT_MODEL_ID, help="LLM Model ID")
    parser.add_argument(
        "--base-url",
        default=DEFAULT_LLM_BASE_URL,
        help="LLM Base URL (for OpenAI compatible providers)",
    )
    parser.add_argument("--api-key", default=DEFAULT_LLM_API_KEY, help="LLM API Key")
    parser.add_argument("--mcp-url", default=DEFAULT_MCP_URL, help="MCP Server URL")
    parser.add_argument(
        "--mcp-config", default=DEFAULT_MCP_CONFIG, help="MCP Server Config"
    )
    parser.add_argument(
        "--skills-directory",
        default=DEFAULT_SKILLS_DIRECTORY,
        help="Directory containing agent skills",
    )

    parser.add_argument(
        "--web",
        action="store_true",
        default=DEFAULT_ENABLE_WEB_UI,
        help="Enable Pydantic AI Web UI",
    )

    parser.add_argument(
        "--insecure",
        action="store_true",
        help="Disable SSL verification for LLM requests (Use with caution)",
    )

    parser.add_argument("--help", action="store_true", help="Show usage")

    args = parser.parse_args()

    if hasattr(args, "help") and args.help:

        usage()

        sys.exit(0)

    if args.debug:
        for handler in logging.root.handlers[:]:
            logging.root.removeHandler(handler)

        logging.basicConfig(
            level=logging.DEBUG,
            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
            handlers=[logging.StreamHandler()],
            force=True,
        )
        logging.getLogger("pydantic_ai").setLevel(logging.DEBUG)
        logging.getLogger("fastmcp").setLevel(logging.DEBUG)
        logging.getLogger("httpcore").setLevel(logging.DEBUG)
        logging.getLogger("httpx").setLevel(logging.DEBUG)
        logger.setLevel(logging.DEBUG)
        logger.debug("Debug mode enabled")

    create_agent_server(
        provider=args.provider,
        model_id=args.model_id,
        base_url=args.base_url,
        api_key=args.api_key,
        mcp_url=args.mcp_url,
        mcp_config=args.mcp_config,
        skills_directory=args.skills_directory,
        debug=args.debug,
        host=args.host,
        port=args.port,
        enable_web_ui=args.web,
        ssl_verify=not args.insecure,
    )


def usage():
    print(
        f"Arr Mcp ({__version__}): Arr Supervisor Agent\n\n"
        "Usage:\n"
        "--host                [ Host to bind the server to ]\n"
        "--port                [ Port to bind the server to ]\n"
        "--debug               [ Debug mode ]\n"
        "--reload              [ Enable auto-reload ]\n"
        "--provider            [ LLM Provider ]\n"
        "--model-id            [ LLM Model ID ]\n"
        "--base-url            [ LLM Base URL (for OpenAI compatible providers) ]\n"
        "--api-key             [ LLM API Key ]\n"
        "--mcp-url             [ MCP Server URL ]\n"
        "--mcp-config          [ MCP Server Config ]\n"
        "--skills-directory    [ Directory containing agent skills ]\n"
        "--web                 [ Enable Pydantic AI Web UI ]\n"
        "\n"
        "Examples:\n"
        "  [Simple]  arr-agent \n"
        '  [Complex] arr-agent --host "value" --port "value" --debug "value" --reload --provider "value" --model-id "value" --base-url "value" --api-key "value" --mcp-url "value" --mcp-config "value" --skills-directory "value" --web\n'
    )


if __name__ == "__main__":
    agent_server()
