"""
CLI command for testing Pydantic AI agents.

Usage:
    rem ask query-agent "Find all documents by Sarah" --model anthropic:claude-sonnet-4-5-20250929
    rem ask schemas/query-agent.yaml "What is the weather?" --temperature 0.7 --max-turns 5
    rem ask my-agent "Hello" --stream --version 1.2.0
"""

import asyncio
import json
import sys
from pathlib import Path
from typing import Any

import click
from loguru import logger

from ...agentic.context import AgentContext
from ...agentic.providers.pydantic_ai import create_agent
from ...agentic.query import AgentQuery
from ...settings import settings
from ...utils.schema_loader import load_agent_schema


async def load_schema_from_registry(
    name: str, version: str | None = None
) -> dict[str, Any]:
    """
    Load agent schema from registry (database or cache).

    TODO: Implement schema registry with:
    - Database table: agent_schemas (name, version, schema_json, created_at)
    - Cache layer: Redis/in-memory for fast lookups
    - Versioning: semantic versioning with latest fallback

    Args:
        name: Schema name (e.g., "query-agent", "rem-agents-query-agent")
        version: Optional version (e.g., "1.2.0", defaults to latest)

    Returns:
        Agent schema as dictionary

    Example:
        schema = await load_schema_from_registry("query-agent", version="1.0.0")
    """
    # TODO: Implement database/cache lookup
    # from ...db import get_db_pool
    # async with get_db_pool() as pool:
    #     if version:
    #         query = "SELECT schema_json FROM agent_schemas WHERE name = $1 AND version = $2"
    #         row = await pool.fetchrow(query, name, version)
    #     else:
    #         query = "SELECT schema_json FROM agent_schemas WHERE name = $1 ORDER BY created_at DESC LIMIT 1"
    #         row = await pool.fetchrow(query, name)
    #
    #     if not row:
    #         raise ValueError(f"Schema not found: {name} (version: {version or 'latest'})")
    #
    #     return json.loads(row["schema_json"])

    raise NotImplementedError(
        f"Schema registry not implemented yet. Please use a file path instead.\n"
        f"Attempted to load: {name} (version: {version or 'latest'})"
    )


async def run_agent_streaming(
    agent,
    prompt: str,
    max_turns: int = 10,
    context: AgentContext | None = None,
    max_iterations: int | None = None,
) -> None:
    """
    Run agent in streaming mode using agent.iter() with usage limits.

    Design Pattern (from carrier):
    - Use agent.iter() for complete execution with tool call visibility
    - run_stream() stops after first output, missing tool calls
    - Stream tool call markers: [Calling: tool_name]
    - Stream text content deltas as they arrive
    - Show final structured result

    Args:
        agent: Pydantic AI agent
        prompt: Complete prompt (includes system context + history + query)
        max_turns: Maximum turns for agent execution (not used in current API)
        context: Optional AgentContext for session persistence
        max_iterations: Maximum iterations/requests (from agent schema or settings)
    """
    from pydantic_ai import UsageLimits
    from rem.utils.date_utils import to_iso_with_z, utc_now

    logger.info("Running agent in streaming mode...")

    try:
        # Import event types for streaming
        from pydantic_ai import Agent as PydanticAgent
        from pydantic_ai.messages import PartStartEvent, PartDeltaEvent, TextPartDelta, ToolCallPart

        # Accumulate assistant response for session persistence
        assistant_response_parts = []

        # Use agent.iter() to get complete execution with tool calls
        usage_limits = UsageLimits(request_limit=max_iterations) if max_iterations else None
        async with agent.iter(prompt, usage_limits=usage_limits) as agent_run:
            async for node in agent_run:
                # Check if this is a model request node (includes tool calls and text)
                if PydanticAgent.is_model_request_node(node):
                    # Stream events from model request
                    request_stream: Any
                    async with node.stream(agent_run.ctx) as request_stream:
                        async for event in request_stream:
                            # Tool call start event
                            if isinstance(event, PartStartEvent) and isinstance(
                                event.part, ToolCallPart
                            ):
                                tool_marker = f"\n[Calling: {event.part.tool_name}]"
                                print(tool_marker, flush=True)
                                assistant_response_parts.append(tool_marker)

                            # Text content delta
                            elif isinstance(event, PartDeltaEvent) and isinstance(
                                event.delta, TextPartDelta
                            ):
                                print(event.delta.content_delta, end="", flush=True)
                                assistant_response_parts.append(event.delta.content_delta)

        print("\n")  # Final newline after streaming

        # Get final result from agent_run
        result = agent_run.result
        if hasattr(result, "output"):
            logger.info("Final structured result:")
            output = result.output
            from rem.agentic.serialization import serialize_agent_result
            output_json = json.dumps(serialize_agent_result(output), indent=2)
            print(output_json)
            assistant_response_parts.append(f"\n{output_json}")

        # Save session messages (if session_id provided and postgres enabled)
        if context and context.session_id and settings.postgres.enabled:
            from ...services.session.compression import SessionMessageStore

            # Extract just the user query from prompt
            # Prompt format from ContextBuilder: system + history + user message
            # We need to extract the last user message
            user_message_content = prompt.split("\n\n")[-1] if "\n\n" in prompt else prompt

            user_message = {
                "role": "user",
                "content": user_message_content,
                "timestamp": to_iso_with_z(utc_now()),
            }

            assistant_message = {
                "role": "assistant",
                "content": "".join(assistant_response_parts),
                "timestamp": to_iso_with_z(utc_now()),
            }

            # Store messages with compression
            store = SessionMessageStore(user_id=context.user_id or settings.test.effective_user_id)
            await store.store_session_messages(
                session_id=context.session_id,
                messages=[user_message, assistant_message],
                user_id=context.user_id,
                compress=True,
            )

            logger.debug(f"Saved conversation to session {context.session_id}")

    except Exception as e:
        logger.error(f"Agent execution failed: {e}")
        raise


async def run_agent_non_streaming(
    agent,
    prompt: str,
    max_turns: int = 10,
    output_file: Path | None = None,
    context: AgentContext | None = None,
    plan: bool = False,
    max_iterations: int | None = None,
) -> dict[str, Any] | None:
    """
    Run agent in non-streaming mode using agent.run() with usage limits.

    Args:
        agent: Pydantic AI agent
        prompt: Complete prompt (includes system context + history + query)
        max_turns: Maximum turns for agent execution (not used in current API)
        output_file: Optional path to save output
        context: Optional AgentContext for session persistence
        plan: If True, output only the generated query (for query-agent)
        max_iterations: Maximum iterations/requests (from agent schema or settings)

    Returns:
        Output data if successful, None otherwise
    """
    from pydantic_ai import UsageLimits
    from rem.utils.date_utils import to_iso_with_z, utc_now

    logger.info("Running agent in non-streaming mode...")

    try:
        # Run agent and get complete result with usage limits
        usage_limits = UsageLimits(request_limit=max_iterations) if max_iterations else None
        result = await agent.run(prompt, usage_limits=usage_limits)

        # Extract output data
        output_data = None
        assistant_content = None
        if hasattr(result, "output"):
            output = result.output
            from rem.agentic.serialization import serialize_agent_result
            output_data = serialize_agent_result(output)

            if plan and isinstance(output_data, dict) and "query" in output_data:
                # Plan mode: Output only the query
                # Use sql formatting if possible or just raw string
                assistant_content = output_data["query"]
                print(assistant_content)
            else:
                # Normal mode
                assistant_content = json.dumps(output_data, indent=2)
                print(assistant_content)
        else:
            # Fallback for text-only results
            assistant_content = str(result)
            print(assistant_content)

        # Save to file if requested
        if output_file and output_data:
            await _save_output_file(output_file, output_data)

        # Save session messages (if session_id provided and postgres enabled)
        if context and context.session_id and settings.postgres.enabled:
            from ...services.session.compression import SessionMessageStore

            # Extract just the user query from prompt
            # Prompt format from ContextBuilder: system + history + user message
            # We need to extract the last user message
            user_message_content = prompt.split("\n\n")[-1] if "\n\n" in prompt else prompt

            user_message = {
                "role": "user",
                "content": user_message_content,
                "timestamp": to_iso_with_z(utc_now()),
            }

            assistant_message = {
                "role": "assistant",
                "content": assistant_content,
                "timestamp": to_iso_with_z(utc_now()),
            }

            # Store messages with compression
            store = SessionMessageStore(user_id=context.user_id or settings.test.effective_user_id)
            await store.store_session_messages(
                session_id=context.session_id,
                messages=[user_message, assistant_message],
                user_id=context.user_id,
                compress=True,
            )

            logger.debug(f"Saved conversation to session {context.session_id}")

        return output_data

    except Exception as e:
        logger.error(f"Agent execution failed: {e}")
        raise


async def _load_input_file(
    file_path: Path, user_id: str | None = None
) -> str:
    """
    Load content from input file using ContentService.

    Simple parse operation - just extracts content without creating Resources.

    Args:
        file_path: Path to input file
        user_id: Optional user ID (not used for simple parse)

    Returns:
        Parsed file content as string (markdown format)
    """
    from ...services.content import ContentService

    # Create ContentService instance
    content_service = ContentService()

    # Parse file (read-only, no database writes)
    logger.info(f"Parsing file: {file_path}")
    result = content_service.process_uri(str(file_path))
    content = result["content"]

    logger.info(
        f"Loaded {len(content)} characters from {file_path.suffix} file using {result['provider']}"
    )
    return content


async def _save_output_file(file_path: Path, data: dict[str, Any]) -> None:
    """
    Save output data to file in YAML format.

    Args:
        file_path: Path to output file
        data: Data to save
    """
    import yaml

    with open(file_path, "w", encoding="utf-8") as f:
        yaml.safe_dump(data, f, default_flow_style=False, allow_unicode=True, sort_keys=False)

    logger.success(f"Output saved to: {file_path}")


@click.command()
@click.argument("name_or_query")
@click.argument("query", required=False)
@click.option(
    "--model",
    "-m",
    default=None,
    help=f"LLM model (default: {settings.llm.default_model})",
)
@click.option(
    "--temperature",
    "-t",
    type=float,
    default=None,
    help=f"Temperature for generation (default: {settings.llm.default_temperature})",
)
@click.option(
    "--max-turns",
    type=int,
    default=10,
    help="Maximum turns for agent execution (default: 10)",
)
@click.option(
    "--version",
    "-v",
    default=None,
    help="Schema version (for registry lookup, defaults to latest)",
)
@click.option(
    "--stream/--no-stream",
    default=False,
    help="Enable streaming mode (default: disabled)",
)
@click.option(
    "--user-id",
    default=None,
    help="User ID for context (default: from settings.test.effective_user_id)",
)
@click.option(
    "--session-id",
    default=None,
    help="Session ID for context (default: auto-generated)",
)
@click.option(
    "--input-file",
    "-i",
    type=click.Path(exists=True, path_type=Path),
    default=None,
    help="Read input from file instead of QUERY argument (supports PDF, TXT, Markdown)",
)
@click.option(
    "--output-file",
    "-o",
    type=click.Path(path_type=Path),
    default=None,
    help="Write output to file (YAML format)",
)
@click.option(
    "--plan",
    is_flag=True,
    default=False,
    help="Output only the generated plan/query (useful for query-agent)",
)
def ask(
    name_or_query: str,
    query: str | None,
    model: str | None,
    temperature: float | None,
    max_turns: int,
    version: str | None,
    stream: bool,
    user_id: str | None,
    session_id: str | None,
    input_file: Path | None,
    output_file: Path | None,
    plan: bool,
):
    """
    Run an agent with a query or file input.

    Arguments:
        NAME_OR_QUERY: Agent schema name OR query string.
        QUERY: Query string (if first arg is agent name).

    Examples:
        # Simple query (uses default 'rem' agent)
        rem ask "What documents did I upload?"

        # Explicit agent
        rem ask contract-analyzer "Analyze this contract"

        # Process file
        rem ask contract-analyzer -i contract.pdf -o output.yaml
    """
    # Smart argument handling
    name = "rem"  # Default agent
    
    if query is None and not input_file:
        # Single argument provided
        # Heuristic: If it looks like a schema file or known agent, treat as name
        # Otherwise treat as query
        if name_or_query.endswith((".yaml", ".yml", ".json")) or name_or_query in ["rem", "query-agent", "rem-query-agent"]:
             # It's an agent name, query is missing (unless input_file)
             name = name_or_query
             # Query remains None, _ask_async will check input_file
        else:
             # It's a query, use default agent
             query = name_or_query
    elif query is not None:
        # Two arguments provided
        name = name_or_query

    # Resolve user_id from settings if not provided
    effective_user_id = user_id or settings.test.effective_user_id

    asyncio.run(
        _ask_async(
            name=name,
            query=query,
            model=model,
            temperature=temperature,
            max_turns=max_turns,
            version=version,
            stream=stream,
            user_id=effective_user_id,
            session_id=session_id,
            input_file=input_file,
            output_file=output_file,
            plan=plan,
        )
    )


async def _ask_async(
    name: str,
    query: str | None,
    model: str | None,
    temperature: float | None,
    max_turns: int,
    version: str | None,
    stream: bool,
    user_id: str,
    session_id: str | None,
    input_file: Path | None,
    output_file: Path | None,
    plan: bool,
):
    """Async implementation of ask command."""
    import uuid
    from ...agentic.context_builder import ContextBuilder

    # Validate input arguments
    if not query and not input_file:
        logger.error("Either QUERY argument or --input-file must be provided")
        sys.exit(1)

    if query and input_file:
        logger.error("Cannot use both QUERY argument and --input-file")
        sys.exit(1)

    # Load input from file if specified
    if input_file:
        logger.info(f"Loading input from file: {input_file}")
        query = await _load_input_file(input_file, user_id=user_id)

    # Load schema using centralized utility
    # Handles both file paths and schema names automatically
    # Falls back to database LOOKUP if not found in filesystem
    logger.info(f"Loading schema: {name} (version: {version or 'latest'})")
    try:
        schema = load_agent_schema(name, user_id=user_id)
    except FileNotFoundError as e:
        logger.error(str(e))
        sys.exit(1)

    # Generate session ID if not provided
    if not session_id:
        session_id = str(uuid.uuid4())
        logger.info(f"Generated session ID: {session_id}")

    # Build context with session history using ContextBuilder
    # This provides:
    # - System context message with date and user profile hints
    # - Compressed session history (if session exists)
    # - Proper message structure for agent
    logger.info(f"Building context for user {user_id}, session {session_id}")

    # Prepare new message for ContextBuilder
    new_messages = [{"role": "user", "content": query}]

    # Build context with session history
    context, messages = await ContextBuilder.build_from_headers(
        headers={
            "X-User-Id": user_id,
            "X-Session-Id": session_id,
        },
        new_messages=new_messages,
    )

    # Override model if specified via CLI flag
    if model:
        context.default_model = model

    logger.info(
        f"Creating agent: model={context.default_model}, stream={stream}, max_turns={max_turns}, messages={len(messages)}"
    )

    # Create agent
    agent = await create_agent(
        context=context,
        agent_schema_override=schema,
        model_override=model,
    )

    # Temperature is now handled in agent factory (schema override or settings default)
    if temperature is not None:
        logger.warning(
            f"CLI temperature override ({temperature}) not yet supported. "
            "Use agent schema 'override_temperature' field or LLM__DEFAULT_TEMPERATURE setting."
        )

    # Combine messages into single prompt
    # ContextBuilder already assembled: system context + history + new message
    prompt = "\n\n".join(msg.content for msg in messages)

    # Run agent with session persistence
    if stream:
        await run_agent_streaming(agent, prompt, max_turns=max_turns, context=context)
    else:
        await run_agent_non_streaming(
            agent,
            prompt,
            max_turns=max_turns,
            output_file=output_file,
            context=context,
            plan=plan,
        )

    # Log session ID for reuse
    logger.success(f"Session ID: {session_id} (use --session-id to continue this conversation)")


def register_command(parent_group):
    """Register ask command with parent CLI group."""
    parent_group.add_command(ask)
