import argparse
import asyncio
import json
import logging
import os
import sys
from typing import Callable, List
from mcp.server.fastmcp import FastMCP, Context

from puli_mcp_server.mcp_server.models import ChangeSet
from puli_mcp_server.proxy_client import ProxyClient
from puli_mcp_server.embedding_client.client import EmbeddingClient
from puli_mcp_server.embedding_client.config import EmbeddingConfig
from puli_mcp_server.llm_agent.llm_agent import LLMAgent
from puli_mcp_server.llm_agent.config import LLMAgentConfig, resolve_llm_provider
from puli_mcp_server.llm_agent.models import LLMQueryRequest, RiskAssessmentList

logger = logging.getLogger(__name__)

mcp = FastMCP("code-reviewer")

# Initialize clients with remote configuration
proxy_client = ProxyClient()

# Resolve LLM provider from BYO API keys (priority: Anthropic > OpenAI > Google)
# This also sets the appropriate env var for pydantic-ai (e.g. ANTHROPIC_API_KEY, OPENAI_API_KEY, GEMINI_API_KEY)
_byo_provider, _byo_model, _byo_key_name = resolve_llm_provider()

# Guest users MUST provide a BYO API key
if proxy_client.role == "GUEST" and not _byo_provider:
    logger.error(
        "Guest users must provide their own LLM API key. "
        "Set one of: BYO_ANTHROPIC_API_KEY, BYO_OPENAI_API_KEY, or BYO_GOOGLE_API_KEY"
    )
    sys.exit(1)

# Resolve OpenAI API key for embeddings and default LLM
_remote_openai_key = proxy_client.mcp_config.get("config", {}).get("OPENAI_API_KEY")
if _byo_provider != "openai" and _remote_openai_key:
    os.environ["OPENAI_API_KEY"] = _remote_openai_key
if not _byo_provider and _remote_openai_key:
    os.environ["OPENAI_API_KEY"] = _remote_openai_key

# Initialize embedding: use local OpenAI client if key is available, otherwise proxy endpoint
_has_openai_key = bool(os.environ.get("OPENAI_API_KEY"))
if _has_openai_key:
    embedding_config = EmbeddingConfig.from_remote(proxy_client.mcp_config["config"])
    _local_embedding_client = EmbeddingClient(config=embedding_config)
    generate_embedding: Callable[[str], List[float]] = _local_embedding_client.generate_embedding
    logger.info("Embedding mode: local (OpenAI API key available)")
else:
    generate_embedding: Callable[[str], List[float]] = proxy_client.generate_embedding
    logger.info("Embedding mode: proxy (no local OpenAI API key)")

llm_config = LLMAgentConfig.from_remote(
    proxy_client.mcp_config["config"],
    proxy_client.mcp_config["prompts"],
    provider_override=_byo_provider,
    model_override=_byo_model,
)
llm_agent = LLMAgent(config=llm_config)


def log_configuration():
    """Log MCP configuration after logging is properly initialized."""
    logger.info("=" * 60)
    logger.info("MCP Configuration pulled from proxy:")
    logger.info("  Config keys: %s", list(proxy_client.mcp_config.get("config", {}).keys()))
    logger.info("  Prompt keys: %s", list(proxy_client.mcp_config.get("prompts", {}).keys()))
    logger.info("=" * 60)
    logger.info("Embedding mode: %s", "local" if _has_openai_key else "proxy")
    logger.info("LLM Config - Provider: %s, Model: %s, Temperature: %s", 
                llm_config.provider, llm_config.model, llm_config.temperature)
    logger.info("System prompt length: %d characters", len(llm_config.system_prompt))
    if _byo_key_name:
        logger.info("LLM API key source: %s (provider: %s)", _byo_key_name, _byo_provider)
    elif _remote_openai_key:
        logger.info("LLM API key source: Remote config (proxy /config/mcp) - OpenAI default")
    else:
        logger.warning("LLM API key: NOT SET - LLM calls will fail")
    logger.info("=" * 60)


# --- 3. Define the Tool ---

@mcp.tool()
async def puli_herd(change_sets: List[ChangeSet], ctx: Context) -> str:
    """
    Puli Guard is a tool that analyzes a list of code changes for security risks, complexity, and
    infrastructure impact by comparing them with historical incidents.
    """

    total_risk_assessments = RiskAssessmentList()

    for i, change_set in enumerate(change_sets, 1):
        await ctx.info(f"Analyzing change set {i}: {change_set.goal}")

        # 1. Create embedding vector
        embedding_text = change_set.to_embedding_string()
        vector = generate_embedding(embedding_text)

        # 2. Query proxy for similar historical incidents and relevant chaos patterns (in parallel)
        similar_incidents, similar_chaos_patterns = await asyncio.gather(
            asyncio.to_thread(proxy_client.search_incidents, query_vector=vector),
            asyncio.to_thread(proxy_client.search_chaos_patterns, query_vector=vector)
        )

        # 3. Create LLM query request and get analysis
        query_request = LLMQueryRequest(
            change_set=change_set,
            historical_incidents=similar_incidents,
            relevant_chaos_patterns=similar_chaos_patterns
        )
        risk_assessments = await llm_agent.query(query_request)
        total_risk_assessments.risk_assessments.extend(risk_assessments)

    return total_risk_assessments.to_str()


class MockContext:
    """Mock context for CLI mode that prints info messages to stdout."""
    async def info(self, msg: str):
        print(f"[INFO] {msg}")


async def run_from_file(file_path: str) -> str:
    """Run analysis from a JSON file containing change sets."""
    with open(file_path) as f:
        data = json.load(f)

    change_sets = [ChangeSet(**cs) for cs in data]
    return await puli_herd(change_sets, MockContext())


def main():
    # Configure logging
    logging.basicConfig(
        level=logging.INFO,
        format='[%(asctime)s] %(levelname)s - %(name)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    
    # Log configuration after logging is set up
    # log_configuration()
    
    parser = argparse.ArgumentParser(description="Puli Code Reviewer")
    parser.add_argument("--file", "-f", help="JSON file with change sets (CLI mode)")
    args = parser.parse_args()

    if args.file:
        result = asyncio.run(run_from_file(args.file))
        print(result)
    else:
        mcp.run()  # Original MCP mode


if __name__ == "__main__":
    main()
