import argparse
import asyncio
import json
import logging
import os
import sys
import time
from typing import Callable, List, Optional
import concurrent.futures

from mcp.server.fastmcp import FastMCP, Context

from puli_mcp_server.mcp_server.models import ChangeSet, CodingSessionAssessmentSummary
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_mcp_server.vertex_token_manager import VertexTokenManager
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: controls LLM-level Vertex usage (BYO key overrides remote Vertex config)
_is_vertex = _remote_llm_provider in _VERTEX_PROVIDERS and not _byo_provider

# _is_vertex_embedding: controls embedding-level Vertex usage (independent of BYO key)
# When remote config points to Vertex, we always generate embeddings via Vertex —
# regardless of whether the customer is using a BYO LLM key.
_is_vertex_embedding = _remote_llm_provider in _VERTEX_PROVIDERS

# Detect whether the customer has supplied their own GCP project/location env vars.
# If so, we use their ADC credentials directly (case 1) instead of the proxy-minted
# SA token (case 2). Both VERTEX_PROJECT and VERTEX_LOCATION (or the embedding-specific
# VERTEX_EMBEDDING_LOCATION) must be present for us to skip the SA token path.
_customer_vertex_project = os.environ.get("VERTEX_PROJECT") or os.environ.get("GOOGLE_CLOUD_PROJECT")
_customer_vertex_location = (
    os.environ.get("VERTEX_EMBEDDING_LOCATION")
    or os.environ.get("VERTEX_LOCATION")
    or os.environ.get("GOOGLE_CLOUD_LOCATION")
)
_customer_has_vertex_env = bool(_customer_vertex_project and _customer_vertex_location)

# --- Vertex AI token manager (SA-based auth) ---
# Needed when remote config is Vertex AND the customer has NOT provided their own
# GCP project/location. In that case the proxy mints a short-lived SA token used
# for both LLM (if _is_vertex) and embeddings (if _is_vertex_embedding).
# If minting fails (e.g. local dev without SA), fall back to ADC (gcloud login).
_vertex_token_manager: Optional[VertexTokenManager] = None
if _is_vertex_embedding and not _customer_has_vertex_env:
    try:
        _vertex_token_manager = VertexTokenManager(proxy_client)
        # Force an initial token fetch to validate the setup
        _vertex_token_manager.get_credentials()
        logger.info("Vertex AI auth: using proxy-minted SA token")
    except Exception as exc:
        logger.info(
            "Vertex AI auth: SA token not available (%s), falling back to ADC (gcloud login)",
            exc,
        )
        _vertex_token_manager = None

# Resolve OpenAI API key for default LLM (non-embedding uses)
# Skip for google-vertex provider — it uses ADC or SA token, not an API key.
_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")

# --- LLM config (sets GOOGLE_CLOUD_PROJECT / GOOGLE_CLOUD_LOCATION for pydantic-ai when using ADC) ---
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, vertex_token_manager=_vertex_token_manager)

# --- Embedding client selection ---
#
# Four cases, in priority order:
#
#   1. Remote config = Vertex AND customer sets VERTEX_PROJECT + VERTEX_LOCATION
#      → embed locally using the customer's own GCP credentials (ADC).
#        User code never touches the proxy; the customer pays for their own project.
#
#   2. Remote config = Vertex AND no customer GCP env vars
#      → embed locally using the proxy-minted SA token (or ADC fallback).
#        This is the standard managed-Vertex path.
#
#   3. Remote config = OpenAI AND customer set BYO_OPENAI_API_KEY
#      → embed locally using the customer's own OpenAI key (already resolved by
#        resolve_llm_provider() into OPENAI_API_KEY).
#        User code never touches the proxy.
#
#   4. Remote config = OpenAI AND no BYO key
#      → proxy handles embeddings via /embeddings/generate.
#
# IMPORTANT: We clear GOOGLE_CLOUD_LOCATION *before* creating any Vertex embedding
# client. _setup_vertex_env (called by LLMAgentConfig.from_remote 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 at construction time, overriding the explicit `location`
# constructor arg. Popping it here ensures the embedding client uses its own location.
if _is_vertex_embedding:
    _saved_gcl = os.environ.pop("GOOGLE_CLOUD_LOCATION", None)
    _remote_config = proxy_client.mcp_config.get("config", {})
    try:
        if _customer_has_vertex_env:
            # Case 1: use the customer's own GCP project via ADC (from_env reads env vars)
            _embedding_client = get_embedding_client(provider="vertex")
            _embedding_mode = "local-vertex (customer ADC)"
        else:
            # Case 2: use proxy-minted SA token (or ADC fallback if SA unavailable)
            _vertex_credentials = _vertex_token_manager.get_credentials() if _vertex_token_manager else None
            _embedding_client = get_embedding_client(
                provider="vertex",
                remote_config=_remote_config,
                credentials=_vertex_credentials,
            )
            _embedding_mode = "local-vertex" + (" (SA token)" if _vertex_token_manager else " (ADC)")
    except (ValueError, Exception) as exc:
        logger.error(
            "Failed to initialize local Vertex embedding client: %s. "
            "Case 1 (customer ADC): ensure VERTEX_PROJECT and VERTEX_LOCATION are set. "
            "Case 2 (SA token): ensure ADC is configured (gcloud auth application-default login) "
            "or SA token is available from the proxy.",
            exc,
        )
        sys.exit(1)
    # Restore GOOGLE_CLOUD_LOCATION for any other code that may need it.
    if _saved_gcl:
        os.environ["GOOGLE_CLOUD_LOCATION"] = _saved_gcl
    generate_embedding: Callable[[str], List[float]] = _embedding_client.generate_embedding
    logger.info(
        "Embedding mode: %s (model=%s, project=%s, location=%s)",
        _embedding_mode,
        _embedding_client.config.model,
        _embedding_client.config.project,
        _embedding_client.config.location,
    )
else:
    if _byo_provider == "openai":
        # Case 3: BYO OpenAI key — embed locally (OPENAI_API_KEY already set above)
        try:
            _embedding_client = get_embedding_client(provider="openai")
        except (ValueError, Exception) as exc:
            logger.error(
                "Failed to initialize local OpenAI embedding client: %s. "
                "Ensure BYO_OPENAI_API_KEY is set correctly.",
                exc,
            )
            sys.exit(1)
        generate_embedding: Callable[[str], List[float]] = _embedding_client.generate_embedding
        _embedding_mode = "local-openai (BYO key)"
        logger.info(
            "Embedding mode: %s (model=%s)",
            _embedding_mode,
            _embedding_client.config.model,
        )
    else:
        # Case 4: no BYO key, no Vertex — proxy handles embeddings
        generate_embedding: Callable[[str], List[float]] = proxy_client.generate_embedding
        _embedding_mode = "proxy"
        logger.info("Embedding mode: proxy")


# --- Tenant profile (RAGFlow profiler chat) ---
# Fetch the tenant profile at startup so it's available for LLM context.
# GUEST and EVALUATOR roles don't have RAGFlow access, so skip for them.
# The call is wrapped in a timeout to avoid blocking startup if RAGFlow is slow.
_PROFILER_TIMEOUT_SECONDS = 30
_tenant_profile: Optional[str] = None
if proxy_client.role not in ("GUEST", "EVALUATOR"):

    def _fetch_tenant_profile() -> Optional[str]:
        result = proxy_client.get_tenant_profile(
            "what can you tell me about this company? "
            "what type of company are they, what do they do? what makes them unique?"
            "If some information is not available, don't mention it, reutrn only the information that is available."
        )
        return result.get("answer")

    try:
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
            future = pool.submit(_fetch_tenant_profile)
            _tenant_profile = future.result(timeout=_PROFILER_TIMEOUT_SECONDS)
        logger.info("Tenant profile loaded (%d chars)", len(_tenant_profile or ""))
    except concurrent.futures.TimeoutError:
        logger.warning(
            "Tenant profile fetch timed out after %ds — continuing without it",
            _PROFILER_TIMEOUT_SECONDS,
        )
    except Exception as exc:
        logger.warning("Failed to load tenant profile from RAGFlow: %s", exc)
else:
    logger.info("Skipping tenant profile fetch (role=%s)", proxy_client.role)


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)
    if _customer_has_vertex_env:
        logger.info(
            "Embedding Vertex env: customer-provided (project=%s, location=%s)",
            _customer_vertex_project,
            _customer_vertex_location,
        )
    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 and _vertex_token_manager:
        logger.info("LLM API key source: Vertex AI SA token (proxy-minted)")
    elif _is_vertex:
        logger.info("LLM API key source: Vertex AI ADC (gcloud login)")
    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),
        )

        # 3. Create LLM query request and get analysis
        query_request = LLMQueryRequest(
            company_profile=_tenant_profile,
            domains=similar_domain_indicators,
            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
        )
        risk_assessments = await llm_agent.query(query_request)
        total_risk_assessments.risk_assessments.extend(risk_assessments)

    return total_risk_assessments.to_str()


@mcp.tool()
async def puli_learn(session: CodingSessionAssessmentSummary, ctx: Context) -> str:
    """
    Submit feedback on a completed coding session to help Puli learn.
    Provide the session context and your assessment of each risk Puli identified,
    including whether it was fixed and whether the analysis was useful.
    """
    await ctx.info("Submitting session feedback to Puli...")
    result = await asyncio.to_thread(proxy_client.submit_report, session)
    return result.get("message", "Feedback submitted successfully")


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 puli_herd analysis from a JSON file containing change sets."""
    try:
        with open(file_path) as f:
            data = json.load(f)
        change_sets = [ChangeSet(**cs) for cs in data]
    except (OSError, json.JSONDecodeError, ValueError) as exc:
        logger.error("Failed to load change sets from %s: %s", file_path, exc)
        return f"Error: could not load change sets — {exc}"

    return await puli_herd(change_sets, MockContext())


async def run_learn_from_file(file_path: str) -> str:
    """Run puli_learn from a JSON file containing a CodingSessionAssessmentSummary."""
    try:
        with open(file_path) as f:
            data = json.load(f)
        session = CodingSessionAssessmentSummary(**data)
    except (OSError, json.JSONDecodeError, ValueError) as exc:
        logger.error("Failed to load session assessment from %s: %s", file_path, exc)
        return f"Error: could not load session assessment — {exc}"

    return await puli_learn(session, 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")
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--file", "-f", help="JSON file with change sets — runs puli_herd (CLI mode)")
    group.add_argument("--learn", "-l", help="JSON file with session assessment — runs puli_learn (CLI mode)")
    args = parser.parse_args()

    if args.file:
        result = asyncio.run(run_from_file(args.file))
        print(result)
    elif args.learn:
        result = asyncio.run(run_learn_from_file(args.learn))
        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")
