import os
import uuid
from typing import Dict

from mcp.server.fastmcp import FastMCP

# Temporal dependencies - wrapped to allow local testing without server
try:
    from temporalio.client import Client
    from workflows import InvoiceWorkflow
    HAS_TEMPORAL = True
except ImportError:
    HAS_TEMPORAL = False


async def _client() -> 'Client':
    """Connect to Temporal server or use mock for local testing."""
    if HAS_TEMPORAL:
        return await Client.connect(os.getenv("TEMPORAL_ADDRESS", "localhost:7233"))
    return None


# Mock storage for local testing
_mock_workflows = {}


mcp = FastMCP("invoice_processor")


@mcp.tool()
async def process_invoice(invoice: Dict) -> Dict[str, str]:
    """Start the InvoiceWorkflow with the given invoice JSON."""
    if HAS_TEMPORAL:
        client = await _client()
        handle = await client.start_workflow(
            InvoiceWorkflow.run,
            invoice,
            id=f"invoice-{uuid.uuid4()}",
            task_queue="invoice-task-queue",
        )
        return {"workflow_id": handle.id, "run_id": handle.result_run_id}
    else:
        # Mock for local testing
        workflow_id = f"invoice-{uuid.uuid4()}"
        run_id = str(uuid.uuid4())
        _mock_workflows[workflow_id] = {
            "run_id": run_id,
            "status": "INITIALIZING",
            "invoice": invoice
        }
        return {"workflow_id": workflow_id, "run_id": run_id}


@mcp.tool()
async def approve_invoice(workflow_id: str, run_id: str) -> str:
    """Signal approval for the invoice workflow."""
    if HAS_TEMPORAL:
        client = await _client()
        handle = client.get_workflow_handle(workflow_id=workflow_id, run_id=run_id)
        await handle.signal("ApproveInvoice")
        return "APPROVED"
    else:
        # Mock for local testing
        if workflow_id in _mock_workflows:
            _mock_workflows[workflow_id]["status"] = "APPROVED"
        return "APPROVED"


@mcp.tool()
async def reject_invoice(workflow_id: str, run_id: str) -> str:
    """Signal rejection for the invoice workflow."""
    if HAS_TEMPORAL:
        client = await _client()
        handle = client.get_workflow_handle(workflow_id=workflow_id, run_id=run_id)
        await handle.signal("RejectInvoice")
        return "REJECTED"
    else:
        # Mock for local testing
        if workflow_id in _mock_workflows:
            _mock_workflows[workflow_id]["status"] = "REJECTED"
        return "REJECTED"


@mcp.tool()
async def invoice_status(workflow_id: str, run_id: str) -> str:
    """Return current status of the workflow."""
    if HAS_TEMPORAL:
        client = await _client()
        handle = client.get_workflow_handle(workflow_id=workflow_id, run_id=run_id)
        desc = await handle.describe()
        status = await handle.query("GetInvoiceStatus")
        return f"Invoice with ID {workflow_id} is currently {status}. " \
               f"Workflow status: {desc.status.name}"
    else:
        # Mock for local testing
        if workflow_id in _mock_workflows:
            status = _mock_workflows[workflow_id]["status"]
            return f"Invoice with ID {workflow_id} is currently {status}. Workflow status: {status}"
        return f"Invoice with ID {workflow_id} not found"


if __name__ == "__main__":
    mcp.run(transport="stdio")