"""Coda MCP Server - Quantum computing tools via Model Context Protocol.

A lightweight MCP server that exposes quantum computing capabilities:
- Circuit transpilation between frameworks (Qiskit, Cirq, PennyLane, etc.)
- Circuit simulation (CPU via Qiskit Aer, GPU via CUDA-Q)
- QPU job submission (IonQ, IQM, Rigetti, etc.)
- Research paper search via Exa

All computation happens on remote Coda infrastructure - this server
is just a thin MCP wrapper around HTTP calls.

Usage:
    coda-mcp                    # Run MCP server
    uvx coda-mcp                # Run via uvx (no install needed)
    python -m coda_mcp.server   # Run as module

Environment:
    CODA_API_URL    API endpoint (default: https://api.coda.conductorquantum.com)
    CODA_API_TOKEN  Authentication token for the API
"""

from typing import Literal

from mcp.server.fastmcp import FastMCP

from coda_mcp import client

mcp = FastMCP("coda")

# Type definitions for tool parameters
QuantumFramework = Literal["qiskit", "cirq", "pennylane", "braket", "pyquil", "cudaq", "openqasm"]
SimulationMethod = Literal["qasm", "statevector", "density_matrix"]
SimulationBackend = Literal["aer", "cudaq", "auto"]
QpuBackend = Literal["simulator", "ionq", "ionq-forte", "ionq-forte-enterprise", "iqm", "iqm-emerald", "rigetti", "aqt"]


# === Quantum Circuit Tools ===


@mcp.tool()
async def transpile(source_code: str, transpile_to: QuantumFramework) -> dict:
    """Transpile quantum code to a target framework.

    Accepts code from any supported framework and converts to the target.
    Supported frameworks: qiskit, cirq, pennylane, braket, pyquil, cudaq, openqasm

    Args:
        source_code: Quantum circuit code in any supported framework
        transpile_to: Target framework to convert to
    """
    return await client.transpile(source_code, transpile_to)


@mcp.tool()
async def simulate(
    code: str,
    method: SimulationMethod = "qasm",
    shots: int = 1024,
    seed_simulator: int | None = None,
    backend: SimulationBackend = "auto",
) -> dict:
    """Simulate a quantum circuit.

    Runs the circuit on a local or GPU simulator.

    Args:
        code: Python code containing the quantum circuit definition
        method: "qasm" for sampling, "statevector" for full state, "density_matrix" for mixed states
        shots: Number of shots for qasm simulation
        seed_simulator: Random seed for reproducibility
        backend: "aer" (CPU), "cudaq" (GPU), or "auto" (picks based on qubit count)
    """
    return await client.simulate(code, method, shots, seed_simulator, backend)


@mcp.tool()
async def to_openqasm3(code: str) -> dict:
    """Convert a quantum circuit to OpenQASM 3.0 format.

    Produces portable OpenQASM that can run on various QPU backends.

    Args:
        code: Python code containing the quantum circuit definition
    """
    return await client.to_openqasm3(code)


@mcp.tool()
async def estimate_resources(code: str) -> dict:
    """Estimate resource requirements for a quantum circuit.

    Analyzes qubit count, depth, and gate counts to determine complexity.

    Args:
        code: Python code containing the quantum circuit definition
    """
    return await client.estimate_resources(code)


@mcp.tool()
async def split_circuit(code: str) -> dict:
    """Split a complex circuit using circuit cutting.

    Automatically finds optimal cuts for distributed execution of large circuits.

    Args:
        code: Python code containing the quantum circuit definition
    """
    return await client.split_circuit(code)


# === QPU Tools ===


async def _get_cost_estimate(backend: str, shots: int) -> int:
    """Fetch cost estimate from API. Returns 0 on failure (fail open for UX)."""
    try:
        result = await client.qpu_estimate_cost(backend, shots)
        cost = result.get("estimated_cost_cents", 0)
        return int(cost) if cost else 0
    except Exception:
        return 0


@mcp.tool()
async def qpu_submit(
    openqasm3: str,
    backend: QpuBackend,
    shots: int = 100,
    confirm_submission: bool = False,
) -> dict:
    """Submit a circuit to a QPU backend.

    BILLABLE ACTION: This submits to real quantum hardware and consumes credits.
    You MUST set confirm_submission=true to proceed.

    IMPORTANT: Before setting confirm_submission=true, you MUST ask the user
    for explicit confirmation. Present them with the backend, shot count, and
    estimated cost, then ask "Do you want to submit this job? (yes/no)".
    Only set confirm_submission=true if the user explicitly agrees.

    Returns immediately with job ID. Results arrive asynchronously.
    Credits are consumed first; remaining cost becomes overage if enabled in user settings.

    Args:
        openqasm3: OpenQASM 3.0 circuit source code
        backend: Target device (simulator, ionq, ionq-forte, iqm, rigetti, aqt)
        shots: Number of circuit executions (1-10000)
        confirm_submission: Must be true to proceed - this action consumes credits
    """
    if not confirm_submission:
        estimated_cost = await _get_cost_estimate(backend, shots)
        return {
            "success": False,
            "error": "QPU submission requires explicit confirmation. Set confirm_submission=true to proceed. This action consumes credits.",
            "requires_confirmation": True,
            "backend": backend,
            "shots": shots,
            "estimated_cost_cents": estimated_cost,
        }
    return await client.qpu_submit(openqasm3, backend, shots, accept_overage=True)


@mcp.tool()
async def qpu_status(job_id: str) -> dict:
    """Check status of a submitted QPU job.

    Args:
        job_id: Job ID from qpu_submit
    """
    return await client.qpu_status(job_id)


@mcp.tool()
async def qpu_devices() -> dict:
    """List available QPU devices.

    Returns device names, IDs, providers, and technology types.
    """
    return await client.qpu_devices()


# === Search Tools ===


@mcp.tool()
async def search_papers(
    query: str,
    search_type: str = "auto",
    num_results: int = 10,
    category: str | None = None,
    include_domains: list[str] | None = None,
    exclude_domains: list[str] | None = None,
) -> dict:
    """Search for papers and research content.

    Uses Exa API for semantic search across academic and technical content.

    Args:
        query: Search query string
        search_type: "auto", "neural", or "keyword"
        num_results: Number of results (max 100)
        category: Optional filter (e.g., "research paper", "news")
        include_domains: Only include results from these domains
        exclude_domains: Exclude results from these domains
    """
    return await client.search_papers(query, search_type, num_results, category, include_domains, exclude_domains)


@mcp.tool()
async def get_paper(
    urls: list[str],
    include_text: bool = True,
    max_characters: int | None = None,
) -> dict:
    """Get full page contents for URLs.

    Fetches text, summaries, and metadata from web pages.

    Args:
        urls: List of URLs to fetch
        include_text: Whether to include full text content
        max_characters: Max characters per page (None for full)
    """
    return await client.get_paper(urls, include_text, max_characters)


def main():
    """Run the MCP server."""
    mcp.run()


if __name__ == "__main__":
    main()
