import argparse
import asyncio
import json
import logging
import os
import sys
from typing import Callable, List
import time

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.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
from puli_embedding_client import get_embedding_client

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 LLM provider from remote config (may be google-vertex or google-vertex-anthropic for EVALUATOR)
_VERTEX_PROVIDERS = {"google-vertex", "google-vertex-anthropic"}
_remote_llm_provider = proxy_client.mcp_config.get("config", {}).get("LLM_PROVIDER")
_is_vertex = _remote_llm_provider in _VERTEX_PROVIDERS and not _byo_provider

# Resolve OpenAI API key for default LLM (non-embedding uses)
# Skip for google-vertex provider — it uses ADC, not an API key.
# When _is_vertex is true, both LLM and embeddings go directly to Vertex AI via ADC.
_remote_openai_key = proxy_client.mcp_config.get("config", {}).get("OPENAI_API_KEY")
if not _is_vertex:
    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
else:
    logger.info("Vertex AI mode: skipping OpenAI API key setup (using ADC)")

# --- LLM config first (sets GOOGLE_CLOUD_PROJECT / GOOGLE_CLOUD_LOCATION for pydantic-ai) ---
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)

# Embedding mode: when Vertex AI is active, generate embeddings locally via ADC
# so that user code/text never touches the proxy. Otherwise, use the proxy endpoint.
# Note: remote config (project/location) is intentionally sourced from the proxy —
# it is the central source of truth for which GCP project embeddings should target.
#
# IMPORTANT: We clear GOOGLE_CLOUD_LOCATION *before* creating the embedding client.
# _setup_vertex_env (called above) sets GOOGLE_CLOUD_LOCATION to the LLM region
# (e.g. "global"), but the embedding client needs a different region (e.g.
# "us-central1"). google.genai.Client reads GOOGLE_CLOUD_LOCATION lazily at
# request time, which would override the explicit `location` constructor arg.
# Removing it ensures the embedding client uses its own configured location.
if _is_vertex:
    _saved_gcl = os.environ.pop("GOOGLE_CLOUD_LOCATION", None)
    _remote_config = proxy_client.mcp_config.get("config", {})
    try:
        _embedding_client = get_embedding_client(
            provider="vertex",
            remote_config=_remote_config,
        )
    except (ValueError, Exception) as exc:
        logger.error(
            "Failed to initialize local Vertex embedding client: %s. "
            "Ensure VERTEX_PROJECT is set in remote config and "
            "ADC is configured (run: gcloud auth application-default login).",
            exc,
        )
        sys.exit(1)
    # Restore GOOGLE_CLOUD_LOCATION for any other code that may need it,
    # but the embedding client now holds its own location internally.
    if _saved_gcl:
        os.environ["GOOGLE_CLOUD_LOCATION"] = _saved_gcl
    generate_embedding: Callable[[str], List[float]] = _embedding_client.generate_embedding
    _embedding_mode = "local-vertex"
    logger.info(
        "Embedding mode: local Vertex AI (model=%s, project=%s, location=%s)",
        _embedding_client.config.model,
        _embedding_client.config.project,
        _embedding_client.config.location,
    )
else:
    generate_embedding: Callable[[str], List[float]] = proxy_client.generate_embedding
    _embedding_mode = "proxy"
    logger.info("Embedding mode: proxy")


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", _embedding_mode)
    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 _is_vertex:
        logger.info("LLM API key source: Vertex AI ADC (google-vertex 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)


async def _run_collection_search(
    search_fn: Callable[..., List],
    vector: List[float],
):
    """Run a single proxy search with bounded concurrency."""
    return await asyncio.to_thread(search_fn, query_vector=vector)


# --- 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 all required collections in parallel (bounded by semaphore).
        (
            similar_incidents,
            similar_chaos_patterns,
            similar_best_practices,
            similar_security_items,
            similar_regulations,
            similar_domain_indicators,
        ) = await asyncio.gather(
            _run_collection_search(proxy_client.search_incidents, vector),
            _run_collection_search(proxy_client.search_chaos_patterns, vector),
            _run_collection_search(proxy_client.search_best_practices, vector),
            _run_collection_search(proxy_client.search_security_items, vector),
            _run_collection_search(proxy_client.search_regulations, vector),
            _run_collection_search(proxy_client.search_domain_indicators, vector),
        )

        #similar_incidents.extend(demo_orianted_incidents)

        # 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,
            relevant_best_practices=similar_best_practices,
            relevant_security_items=similar_security_items,
            relevant_regulations=similar_regulations,
            relevant_domain_indicators=similar_domain_indicators,
        )
        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__":
    _start = time.perf_counter()
    main()
    _elapsed = time.perf_counter() - _start
    print(f"\n[Timer] main() completed in {_elapsed:.3f}s")
