"""Durable MCP calculation submission and owner-scoped result retrieval.

Only safe server-owned binding fields are persisted. OAuth tokens and context
handles are never used to resume a calculation.
"""
from datetime import datetime, timezone
from uuid import uuid5, NAMESPACE_URL
from fastapi import HTTPException
from app.mcp.models import BoundContext, QueryResult, PageInfo
from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository

QUEUED_TOOLS = frozenset({
    "run_burst_location", "run_leakage_identification", "run_burst_detection",
    "run_simulation", "run_burst_simulation", "run_contaminant_simulation",
    "analyze_valve_closure", "analyze_flushing", "analyze_pressure_regulation",
    "optimize_sensor_placement", "save_simulation_scenario", "create_sensor_placement",
})
CONTROL_TOOLS = frozenset({"get_calculation_job", "get_calculation_result", "cancel_calculation_job", "list_calculation_jobs"})
BINDING_FIELDS = {"project_id", "subject", "client_id", "native_session_id", "purpose", "parent_session_id", "task_id", "run_id", "execution_request_id", "tool_cap", "permission_cap", "task_scope_cap"}


def durable_binding(binding: BoundContext) -> dict:
    return binding.model_dump(mode="json", include=BINDING_FIELDS)


def restore_binding(document: dict) -> BoundContext:
    # Expiration governs admission via a context handle; durable work has its own
    # lifecycle and must revalidate identity, membership and run on every use.
    return BoundContext.model_validate({**document, "context_handle": "durable-calculation", "expires_at": datetime.max.replace(tzinfo=timezone.utc)})


def should_queue(tool, args) -> bool:
    return tool.name in QUEUED_TOOLS and (not tool.mutating or not args.dry_run)


def envelope(source: str, data, *, job_id="list", page=None) -> QueryResult:
    return QueryResult(source=source, source_ref="calculation:" + str(job_id), queried_at=datetime.now(timezone.utc), semantics="task", data=data, page=page, truncated=bool(page and (page.offset or page.has_more)),
        data_gaps=["Calculation status and queue receipts are not analytical evidence."])


def public_job(row: dict) -> dict:
    fields = {"job_id", "status", "tool", "created_at", "started_at", "finished_at", "updated_at", "error", "queue_expires_at", "result_expires_at"}
    from fastapi.encoders import jsonable_encoder
    expiry = row.get("result_expires_at")
    available = row.get("status") == "succeeded" and (expiry is None or expiry > datetime.now(timezone.utc))
    return {**jsonable_encoder({key: row[key] for key in fields if key in row}),
            "result_available": available, "poll_after_ms": 1500}


async def submit_calculation(scope, binding, tool, args, call_id: str, *, job_timeout_seconds: int = 3600) -> QueryResult:
    from app.mcp.query_service import authorize_query
    await authorize_query(scope, binding, tool, args)
    if binding.purpose == "task_reader":
        raise HTTPException(403, "Read-only task agents cannot start calculations")
    if tool.mutating:
        from app.mcp.mutation_service import (
            preview_change, check_digest, mutation_arguments, retryable_capacity_rejection,
        )
        from app.infra.db.postgresql.mcp_queries import McpQueryRepository
        mutation_repo = McpQueryRepository(scope.connection, scope.project.project_id)
        mutation_key = "mutation:" + binding.subject + ":" + binding.native_session_id
        approved_parameters = {"tool": tool.name, "arguments": mutation_arguments(args, preview=False)}
        async with scope.connection.transaction():
            await mutation_repo.lock_receipt(mutation_key, str(args.request_id))
            previous = await mutation_repo.receipt(mutation_key, str(args.request_id))
            if previous:
                if previous["parameters"] != approved_parameters:
                    raise HTTPException(409, "request_id already has different parameters")
                if previous["result"].get("status") == "completed":
                    return QueryResult.model_validate(previous["result"]["output"])
                if not retryable_capacity_rejection(previous["result"]):
                    raise HTTPException(409, "Operation already started; inspect the existing calculation or mutation status")
        check_digest(args, await preview_change(scope, binding, tool.name, args))
    parameters = args.model_dump(mode="json", exclude={"context_handle", "timeout_seconds"})
    request_id = str(args.request_id) if tool.mutating else call_id
    key = str(uuid5(NAMESPACE_URL, ":".join([binding.subject, binding.client_id, binding.native_session_id, str(binding.task_id), str(binding.run_id), request_id])))
    repo = CalculationJobRepository(scope.connection, scope.project.project_id)
    if binding.task_id:
        from app.mcp.query_service import task_service
        from app.mcp.contexts import verify_task
        task_repo = task_service(scope).repo
        # Share the operational transaction lock with finalization. Admission
        # cannot race a pending-job check and then publish work into a closed run.
        async with task_repo.transaction():
            await task_repo.lock(f"task:{binding.task_id}")
            await task_repo.lock(f"run:{binding.run_id}")
            await verify_task(scope, binding)
            row = await repo.enqueue(
                owner_user_id=scope.project.user_id, request_key=key, tool=tool.name,
                parameters=parameters, binding=durable_binding(binding), timeout_seconds=job_timeout_seconds,
            )
    else:
        row = await repo.enqueue(
            owner_user_id=scope.project.user_id, request_key=key, tool=tool.name,
            parameters=parameters, binding=durable_binding(binding), timeout_seconds=job_timeout_seconds,
        )
    await scope.connection.commit()
    return envelope(tool.name, public_job(row), job_id=row["job_id"])


def check_binding(binding, row):
    original = row["binding"]
    if any(str(original.get(key)) != str(getattr(binding, key)) for key in ("subject", "client_id", "native_session_id", "task_id", "run_id", "execution_request_id")):
        raise HTTPException(403, "Calculation belongs to a different session or task run")


async def authorize_original(scope, row):
    from app.mcp.catalog import TOOLS
    from app.mcp.query_service import authorize_query
    from app.mcp.contexts import verify_task
    original = restore_binding(row["binding"])
    tool = TOOLS[row["tool"]]
    # Legacy queue rows stored the old HTTP timeout. It does not govern durable execution.
    parameters = {key: value for key, value in row["parameters"].items() if key != "timeout_seconds"}
    args = tool.model.model_validate({**parameters, "context_handle": "durable-calculation-context-handle"})
    if (str(scope.project.user_id) != str(row["owner_user_id"])
            or str(scope.project.project_id) != str(original.project_id)):
        raise HTTPException(403, "Calculation owner is no longer authorized")
    if original.task_id:
        await verify_task(scope, original)
    await authorize_query(scope, original, tool, args)
    return original, tool, args


async def read_calculation_job(scope, binding, tool, args) -> QueryResult:
    repo = CalculationJobRepository(scope.connection, scope.project.project_id)
    if tool.name == "list_calculation_jobs":
        identity_scope = binding.model_dump(mode="json", include={
            "subject", "client_id", "native_session_id", "task_id", "run_id", "execution_request_id",
        })
        rows = await repo.list(owner_user_id=scope.project.user_id, binding=identity_scope, limit=args.limit + 1, offset=args.offset, status=args.status)
        return envelope(tool.name, [public_job(row) for row in rows[:args.limit]], page=PageInfo(limit=args.limit, offset=args.offset, has_more=len(rows) > args.limit))
    row = await repo.get(args.job_id, owner_user_id=scope.project.user_id,
                         include_result=tool.name == "get_calculation_result")
    if row is None:
        raise HTTPException(404, "Calculation was not found")
    check_binding(binding, row)
    await authorize_original(scope, row)
    if tool.name == "cancel_calculation_job":
        if binding.purpose == "task_reader":
            raise HTTPException(403, "Read-only task agents cannot cancel calculations")
        row = await repo.cancel(args.job_id, owner_user_id=scope.project.user_id)
        await scope.connection.commit()
    if tool.name != "get_calculation_result" or row["status"] != "succeeded":
        return envelope(tool.name, public_job(row), job_id=row["job_id"])
    expiry = row.get("result_expires_at")
    if row.get("result") is None or (expiry is not None and expiry <= datetime.now(timezone.utc)):
        raise HTTPException(410, "Calculation result retention period has expired")
    result = QueryResult.model_validate(row["result"])
    result = await attach_calculation_evidence(scope, row, result)
    value = result.data
    for part in args.data_path:
        if isinstance(value, dict) and part in value:
            value = value[part]
        else:
            raise HTTPException(422, "data_path must select an existing object field")
    if isinstance(value, list):
        result.data = value[args.offset:args.offset + args.limit]
        result.page = PageInfo(limit=args.limit, offset=args.offset, total=len(value), has_more=args.offset + args.limit < len(value))
        result.truncated = result.truncated or args.offset > 0 or result.page.has_more
    elif args.offset:
        raise HTTPException(422, "offset requires a collection data_path")
    else:
        result.data = value
    if len(result.model_dump_json().encode()) > 11000:
        raise HTTPException(413, "Select a narrower data_path or smaller result page")
    return result


async def execute_calculation_job(scope, row) -> QueryResult:
    """Worker entry point; success is persisted by the fenced supervisor."""
    from app.mcp.query_service import execute_query
    original, tool, args = await authorize_original(scope, row)
    result = await execute_query(scope, original, tool, args)
    return result


async def attach_calculation_evidence(scope, row, result):
    from app.mcp.query_service import task_service
    from app.mcp.contexts import verify_task
    original, tool, _ = await authorize_original(scope, row)
    if original.task_id:
        from app.domain.schemas.operational_task import TaskDataSnapshot
        from app.infra.db.postgresql.mcp_queries import McpQueryRepository
        repo = McpQueryRepository(scope.connection, scope.project.project_id)
        key, call_id = "calculation:" + str(row["job_id"]), "result"
        # Result lookup can leave an implicit SELECT transaction open. End it so
        # evidence and its receipt commit together before returning a response.
        await scope.connection.commit()
        task_repo = task_service(scope).repo
        # Match operational finalization's lock order: event lock, then task.
        async with task_repo.transaction():
            await repo.lock_receipt(key, call_id)
            previous = await repo.receipt(key, call_id)
            if previous:
                return QueryResult.model_validate(previous["result"])
            await task_repo.lock(f"task:{original.task_id}")
            await verify_task(scope, original)
            result.source_ref = key
            evidence, _ = await task_service(scope).record_data_snapshot(original.task_id, TaskDataSnapshot(
                source=tool.name, source_ref=key, queried_at=result.queried_at, origin=result.origin,
                data=result.model_dump(mode="json"), truncated=result.truncated))
            result.evidence_id = evidence.evidence_id
            await repo.save_receipt(key, call_id, {"job_id": str(row["job_id"])}, result.model_dump(mode="json"))
    return result