#!/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 pydantic_ai_skills import SkillsToolset
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.13"

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_CUSTOM_SKILLS_DIRECTORY = os.getenv("CUSTOM_SKILLS_DIRECTORY", None)
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, a highly capable coordinator for a suite of specialized media management services.\n"
        "Your role is to analyze user requests and delegate them to the appropriate specialist agent using the available 'ask_..._agent' tools.\n\n"
        "You have access to the following specialists:\n"
        "- **Lidarr Agent (Music)**: Expert in music collection management. Use this for searching artists, monitoring albums, managing music library files, and handling music-related quality profiles.\n"
        "- **Sonarr Agent (TV Shows)**: Specialist in television series management. Use this for tracking shows, monitoring episodes/seasons, viewing release calendars, and managing TV media files.\n"
        "- **Radarr Agent (Movies)**: Focused on movie collection management. Use this for discovering movies, monitoring releases, managing movie files, and handling film collections.\n"
        "- **Prowlarr Agent (Indexers & Search)**: Manages search indexers and trackers. Use this for searching across multiple sites, configuring indexers/trackers, and managing download clients.\n"
        "- **Chaptarr Agent (Books)**: Dedicated to book and author management. Use this for tracking authors, monitoring book releases (ebooks/audiobooks), and managing book libraries.\n"
        "- **Seerr Agent (Discovery & Requests)**: Handles media discovery and user requests. Use this for searching trending/popular media via TMDB, making requests for movies or TV shows, and managing those requests.\n"
        "- **Bazarr Agent (Subtitles)**: Specialist in subtitle management. Use this for searching, downloading, and managing subtitles for movies and TV series already in your library.\n\n"
        "GUIDELINES:\n"
        "1. Identify the core intent of the user's request.\n"
        "2. Delegate to the most appropriate agent(s) using the 'ask_..._agent' tools. You can call multiple agents if required (e.g., check Prowlarr indexers after adding a movie in Radarr).\n"
        "3. Synthesize the results from all tools into a single, cohesive, and helpful response.\n"
        "4. Be professional, warm, and proactive. If an action fails, explain why and suggest alternatives."
    ),
)


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,
    custom_skills_directory: Optional[str] = DEFAULT_CUSTOM_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=provider,
            model_id=model_id,
            base_url=base_url,
            api_key=api_key,
            mcp_url=mcp_url,
            mcp_config=mcp_config,
            custom_skills_directory=custom_skills_directory,
            ssl_verify=ssl_verify,
        )
        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=provider,
            model_id=model_id,
            base_url=base_url,
            api_key=api_key,
            mcp_url=mcp_url,
            mcp_config=mcp_config,
            custom_skills_directory=custom_skills_directory,
            ssl_verify=ssl_verify,
        )
        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=provider,
            model_id=model_id,
            base_url=base_url,
            api_key=api_key,
            mcp_url=mcp_url,
            mcp_config=mcp_config,
            custom_skills_directory=custom_skills_directory,
            ssl_verify=ssl_verify,
        )
        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=provider,
            model_id=model_id,
            base_url=base_url,
            api_key=api_key,
            mcp_url=mcp_url,
            mcp_config=mcp_config,
            custom_skills_directory=custom_skills_directory,
            ssl_verify=ssl_verify,
        )
        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=provider,
            model_id=model_id,
            base_url=base_url,
            api_key=api_key,
            mcp_url=mcp_url,
            mcp_config=mcp_config,
            custom_skills_directory=custom_skills_directory,
            ssl_verify=ssl_verify,
        )
        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=provider,
            model_id=model_id,
            base_url=base_url,
            api_key=api_key,
            mcp_url=mcp_url,
            mcp_config=mcp_config,
            custom_skills_directory=custom_skills_directory,
            ssl_verify=ssl_verify,
        )
        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=provider,
            model_id=model_id,
            base_url=base_url,
            api_key=api_key,
            mcp_url=mcp_url,
            mcp_config=mcp_config,
            custom_skills_directory=custom_skills_directory,
            ssl_verify=ssl_verify,
        )
        logger.info("Bazarr Agent initialized.")
    except Exception as e:
        logger.error(f"Failed to initialize Bazarr Agent: {e}")

    supervisor_skills = []
    supervisor_skills_directories = [get_skills_path()]
    arr_tags = ["lidarr", "sonarr", "radarr", "prowlarr", "chaptarr", "seerr", "bazarr"]
    for tag in arr_tags:
        skill_dir_name = f"arr-mcp-{tag}"

        # Check custom skills directory
        if custom_skills_directory:
            skill_dir_path = os.path.join(custom_skills_directory, skill_dir_name)
            if os.path.exists(skill_dir_path):
                supervisor_skills_directories.append(skill_dir_path)

        # Check default skills directory
        default_skill_path = os.path.join(get_skills_path(), skill_dir_name)
        if os.path.exists(default_skill_path):
            supervisor_skills_directories.append(default_skill_path)

    if custom_skills_directory:
        supervisor_skills_directories.append(custom_skills_directory)

    supervisor_skills.append(SkillsToolset(directories=supervisor_skills_directories))
    logger.info(f"Loaded supervisor skills from: {supervisor_skills_directories}")

    supervisor = Agent(
        name=AGENT_NAME,
        system_prompt=SUPERVISOR_SYSTEM_PROMPT,
        model=model,
        model_settings=settings,
        toolsets=supervisor_skills,
        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."
        try:
            return (
                await child_agents["lidarr"].run(task, usage=ctx.usage, deps=ctx.deps)
            ).output
        except Exception as e:
            logger.error(f"Error in lidarr agent: {e}", exc_info=True)
            return f"Error executing task: {e}"

    @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."
        try:
            return (
                await child_agents["sonarr"].run(task, usage=ctx.usage, deps=ctx.deps)
            ).output
        except Exception as e:
            logger.error(f"Error in sonarr agent: {e}", exc_info=True)
            return f"Error executing task: {e}"

    @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."
        try:
            return (
                await child_agents["radarr"].run(task, usage=ctx.usage, deps=ctx.deps)
            ).output
        except Exception as e:
            logger.error(f"Error in radarr agent: {e}", exc_info=True)
            return f"Error executing task: {e}"

    @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."
        try:
            return (
                await child_agents["prowlarr"].run(task, usage=ctx.usage, deps=ctx.deps)
            ).output
        except Exception as e:
            logger.error(f"Error in prowlarr agent: {e}", exc_info=True)
            return f"Error executing task: {e}"

    @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."
        try:
            return (
                await child_agents["chaptarr"].run(task, usage=ctx.usage, deps=ctx.deps)
            ).output
        except Exception as e:
            logger.error(f"Error in chaptarr agent: {e}", exc_info=True)
            return f"Error executing task: {e}"

    @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."
        try:
            return (
                await child_agents["seerr"].run(task, usage=ctx.usage, deps=ctx.deps)
            ).output
        except Exception as e:
            logger.error(f"Error in seerr agent: {e}", exc_info=True)
            return f"Error executing task: {e}"

    @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."
        try:
            return (
                await child_agents["bazarr"].run(task, usage=ctx.usage, deps=ctx.deps)
            ).output
        except Exception as e:
            logger.error(f"Error in bazarr agent: {e}", exc_info=True)
            return f"Error executing task: {e}"

    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,
    custom_skills_directory: Optional[str] = DEFAULT_CUSTOM_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,
        custom_skills_directory=custom_skills_directory,
        ssl_verify=ssl_verify,
    )

    skills = []
    if custom_skills_directory and os.path.exists(custom_skills_directory):
        skills = load_skills_from_directory(custom_skills_directory)
        logger.info(f"Loaded {len(skills)} skills from {custom_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(
        "--custom-skills-directory",
        default=DEFAULT_CUSTOM_SKILLS_DIRECTORY,
        help="Directory containing additional custom 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:

        parser.print_help()

        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,
        custom_skills_directory=args.custom_skills_directory,
        debug=args.debug,
        host=args.host,
        port=args.port,
        enable_web_ui=args.web,
        ssl_verify=not args.insecure,
    )


if __name__ == "__main__":
    agent_server()
