import argparse
import asyncio
import json
import logging
import os
from typing import 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
from puli_mcp_server.llm_agent.models import LLMQueryRequest

logger = logging.getLogger(__name__)

mcp = FastMCP("code-reviewer")

# Initialize clients with remote configuration
proxy_client = ProxyClient()

# Resolve OpenAI API key: BYO env var takes priority over remote config
_byo_key = os.environ.get("BYO_OPENAI_API_KEY")
_remote_key = proxy_client.mcp_config.get("config", {}).get("OPENAI_API_KEY")
if _byo_key:
    os.environ["OPENAI_API_KEY"] = _byo_key
elif _remote_key:
    os.environ["OPENAI_API_KEY"] = _remote_key

# Initialize embedding client with remote config
embedding_config = EmbeddingConfig.from_remote(proxy_client.mcp_config["config"])
embedding_client = EmbeddingClient(config=embedding_config)

# Initialize LLM agent with remote config and prompts
llm_config = LLMAgentConfig.from_remote(
    proxy_client.mcp_config["config"],
    proxy_client.mcp_config["prompts"]
)
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 Config - Model: %s", embedding_config.model)
    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:
        logger.info("OpenAI API key source: BYO_OPENAI_API_KEY env var")
    elif _remote_key:
        logger.info("OpenAI API key source: Remote config (proxy /config/mcp)")
    else:
        logger.warning("OpenAI API key: NOT SET - LLM and embedding 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.
    """

    results = []

    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 = embedding_client.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_assessment = await llm_agent.query(query_request)
        results.append(risk_assessment.to_str())

    return "\n\n---\n\n".join(results)


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()
