"""FastAPI sub-app for tag endpoints (Airflow 3.x plugin).

Mounted on the Airflow api-server via AirflowPlugin.fastapi_apps at
url_prefix="/api/v1/blocks". Full paths:
  PATCH  /api/v1/blocks/{chunk_id}/tags/status
  PATCH  /api/v1/blocks/{chunk_id}/tags/bulk-status
  PATCH  /api/v1/blocks/{chunk_id}/tags/{tag_id}/name
  DELETE /api/v1/blocks/{chunk_id}/tags/{tag_id}

NOTE: This module runs inside the Airflow api-server process.
Do NOT import from src.*.
"""

import json
import logging
import time

from fastapi import FastAPI, Request, Response
from pydantic import BaseModel, Field, field_validator
from sqlalchemy import text
from typing import Literal

from api.config import get_docstore_namespace
from api.metrics import record_outcome, tag_endpoint_duration_seconds, tag_op_total
from api.middleware import (
    CorrelationIdMiddleware,
    HMACVerificationMiddleware,
    RequestIdMiddleware,
)
from api.milvus import project_verified_tag_names, upsert_chunk_field
from api.responses import error_response, success_response
from api.rls import rls_connection

log = logging.getLogger(__name__)

tags_app = FastAPI(title="TextIQ Tags API")

tags_app.add_middleware(HMACVerificationMiddleware)
tags_app.add_middleware(CorrelationIdMiddleware)
tags_app.add_middleware(RequestIdMiddleware)

NAMESPACE = get_docstore_namespace()

# ---------------------------------------------------------------------------
# Pydantic models
# ---------------------------------------------------------------------------


class TagStatusUpdate(BaseModel):
    tag_id: str
    status: Literal["pending", "verified", "rejected"]


class BulkTagStatusUpdate(BaseModel):
    updates: list[TagStatusUpdate] = Field(min_length=1)


class TagRename(BaseModel):
    name: str = Field(min_length=1, max_length=255)

    @field_validator("name")
    @classmethod
    def must_have_namespace(cls, v: str) -> str:
        if ":" not in v:
            raise ValueError("Tag name must match namespace:value pattern (e.g. domain:nursing)")
        return v


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


class _EarlyReturn(Exception):
    """Signal to exit the rls_connection block and return an error response."""

    def __init__(self, response):
        self.response = response


def _read_chunk_for_update(conn, chunk_id: str) -> dict | None:
    """SELECT ... FOR UPDATE a chunk row and return its value dict, or None."""
    row = conn.execute(
        text(
            "SELECT value FROM data_hierarchical_nodes "
            "WHERE key = :chunk_id AND namespace = :ns FOR UPDATE"
        ),
        {"chunk_id": chunk_id, "ns": NAMESPACE},
    ).fetchone()
    if row is None:
        return None
    return json.loads(row[0]) if isinstance(row[0], str) else row[0]


def _get_tags(value: dict) -> list[dict]:
    """Extract tags list from the chunk value."""
    return value.get("__data__", {}).get("metadata", {}).get("tags", [])


def _write_tags(conn, chunk_id: str, updated_tags: list[dict]) -> None:
    """Write updated tags back to the chunk row."""
    conn.execute(
        text(
            "UPDATE data_hierarchical_nodes "
            "SET value = jsonb_set(value, '{__data__,metadata,tags}', CAST(:tags AS jsonb)) "
            "WHERE key = :chunk_id AND namespace = :ns"
        ),
        {"tags": json.dumps(updated_tags), "chunk_id": chunk_id, "ns": NAMESPACE},
    )


def _make_history_entry(action: str, user_id: str, **extra) -> dict:
    """Build a tag history entry."""
    entry = {"action": action, "by": user_id, "ts": int(time.time())}
    entry.update(extra)
    return entry


# Tag status state machine. Keys = from_status, values = allowed to_status.
# Same-status is excluded → 400 INVALID_TRANSITION at PATCH time.
VALID_TRANSITIONS = {
    "pending": {"verified", "rejected"},
    "verified": {"pending", "rejected"},
    "rejected": {"pending", "verified"},
}

# Maps PATCH target status onto TagHistoryAction enum values
# (suggested | confirmed | rejected | edited | migrated). Re-declared locally
# because this module must not import from src.* (runs in Airflow api-server).
_STATUS_TO_ACTION = {"verified": "confirmed", "rejected": "rejected", "pending": "edited"}


def _status_to_history_action(to_status: str) -> str:
    return _STATUS_TO_ACTION[to_status]


def _check_transition(from_status: str, to_status: str) -> tuple[bool, str | None]:
    """Return (ok, reason). Trusts from_status is one of the three live values."""
    if to_status in VALID_TRANSITIONS[from_status]:
        return True, None
    if to_status == from_status:
        return False, f"Tag already in status '{from_status}'"
    return False, f"Cannot transition from '{from_status}' to '{to_status}'"


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------


@tags_app.patch("/{chunk_id}/tags/status")
def update_tag_status(chunk_id: str, body: TagStatusUpdate, request: Request):
    """Update a single tag's status on a chunk."""
    ctx = request.state.hmac_ctx
    tenant_id = ctx["tenant_id"]
    user_id = ctx["user_id"]
    _t0 = time.perf_counter()
    _op = "update_status"

    try:
        with rls_connection(tenant_id) as conn:
            value = _read_chunk_for_update(conn, chunk_id)
            if value is None:
                raise _EarlyReturn(
                    error_response("NOT_FOUND", "Chunk not found", status_code=404)
                )

            tags = _get_tags(value)
            target = next((t for t in tags if t.get("tag_id") == body.tag_id), None)
            if target is None:
                raise _EarlyReturn(
                    error_response("NOT_FOUND", f"Tag {body.tag_id} not found", status_code=404)
                )

            ok, reason = _check_transition(target["status"], body.status)
            if not ok:
                raise _EarlyReturn(
                    error_response("INVALID_TRANSITION", reason, status_code=400)
                )

            target["status"] = body.status
            target.setdefault("history", []).append(
                _make_history_entry(_status_to_history_action(body.status), user_id)
            )

            _write_tags(conn, chunk_id, tags)
            upsert_chunk_field(
                chunk_id,
                tags=tags,
                verified_tag_names=project_verified_tag_names(tags, chunk_id),
            )

        resp = success_response({"chunk_id": chunk_id, "tag_id": body.tag_id, "status": body.status})
        record_outcome(tag_op_total, resp.status_code, operation=_op, tenant_id=tenant_id)
        tag_endpoint_duration_seconds.labels(operation=_op).observe(time.perf_counter() - _t0)
        return resp

    except _EarlyReturn as e:
        record_outcome(tag_op_total, e.response.status_code, operation=_op, tenant_id=tenant_id)
        tag_endpoint_duration_seconds.labels(operation=_op).observe(time.perf_counter() - _t0)
        return e.response
    except Exception as e:
        log.exception("Tag status update failed: %s", e)
        resp = error_response("INTERNAL_ERROR", "Tag status update failed", status_code=500)
        record_outcome(tag_op_total, resp.status_code, operation=_op, tenant_id=tenant_id)
        tag_endpoint_duration_seconds.labels(operation=_op).observe(time.perf_counter() - _t0)
        return resp


@tags_app.patch("/{chunk_id}/tags/bulk-status")
def bulk_update_tag_status(chunk_id: str, body: BulkTagStatusUpdate, request: Request):
    """Bulk update tag statuses on a chunk. All-or-nothing: if any tag_id is missing, none are updated."""
    ctx = request.state.hmac_ctx
    tenant_id = ctx["tenant_id"]
    user_id = ctx["user_id"]
    _t0 = time.perf_counter()
    _op = "bulk_update_status"

    try:
        with rls_connection(tenant_id) as conn:
            value = _read_chunk_for_update(conn, chunk_id)
            if value is None:
                raise _EarlyReturn(
                    error_response("NOT_FOUND", "Chunk not found", status_code=404)
                )

            tags = _get_tags(value)
            tag_map = {t.get("tag_id"): t for t in tags}

            # Validate all tag_ids exist first
            missing = [u.tag_id for u in body.updates if u.tag_id not in tag_map]
            if missing:
                raise _EarlyReturn(
                    error_response(
                        "NOT_FOUND",
                        f"Tag(s) not found: {', '.join(missing)}",
                        status_code=400,
                    )
                )

            # Reject duplicate tag_ids in the request: the apply loop would
            # see already-mutated state on the second occurrence, bypassing
            # the transition guard.
            seen: set[str] = set()
            duplicates: list[str] = []
            for update in body.updates:
                if update.tag_id in seen:
                    duplicates.append(update.tag_id)
                seen.add(update.tag_id)
            if duplicates:
                raise _EarlyReturn(
                    error_response(
                        "DUPLICATE_TAG_ID",
                        f"Duplicate tag_id(s) in updates: {', '.join(sorted(set(duplicates)))}",
                        status_code=400,
                    )
                )

            # Pre-validate all transitions (all-or-nothing)
            invalid: list[str] = []
            for update in body.updates:
                tag = tag_map[update.tag_id]
                ok, reason = _check_transition(tag["status"], update.status)
                if not ok:
                    invalid.append(f"{update.tag_id} ({reason})")
            if invalid:
                raise _EarlyReturn(
                    error_response(
                        "INVALID_TRANSITION",
                        f"Invalid transition(s): {'; '.join(invalid)}",
                        status_code=400,
                    )
                )

            # Apply updates
            for update in body.updates:
                tag = tag_map[update.tag_id]
                tag["status"] = update.status
                tag.setdefault("history", []).append(
                    _make_history_entry(_status_to_history_action(update.status), user_id)
                )

            _write_tags(conn, chunk_id, tags)
            upsert_chunk_field(
                chunk_id,
                tags=tags,
                verified_tag_names=project_verified_tag_names(tags, chunk_id),
            )

        resp = success_response({
            "chunk_id": chunk_id,
            "updated": [{"tag_id": u.tag_id, "status": u.status} for u in body.updates],
        })
        record_outcome(tag_op_total, resp.status_code, operation=_op, tenant_id=tenant_id)
        tag_endpoint_duration_seconds.labels(operation=_op).observe(time.perf_counter() - _t0)
        return resp

    except _EarlyReturn as e:
        record_outcome(tag_op_total, e.response.status_code, operation=_op, tenant_id=tenant_id)
        tag_endpoint_duration_seconds.labels(operation=_op).observe(time.perf_counter() - _t0)
        return e.response
    except Exception as e:
        log.exception("Bulk tag status update failed: %s", e)
        resp = error_response("INTERNAL_ERROR", "Bulk tag status update failed", status_code=500)
        record_outcome(tag_op_total, resp.status_code, operation=_op, tenant_id=tenant_id)
        tag_endpoint_duration_seconds.labels(operation=_op).observe(time.perf_counter() - _t0)
        return resp


@tags_app.patch("/{chunk_id}/tags/{tag_id}/name")
def rename_tag(chunk_id: str, tag_id: str, body: TagRename, request: Request):
    """Rename a tag on a chunk."""
    ctx = request.state.hmac_ctx
    tenant_id = ctx["tenant_id"]
    user_id = ctx["user_id"]
    _t0 = time.perf_counter()
    _op = "rename"

    try:
        with rls_connection(tenant_id) as conn:
            value = _read_chunk_for_update(conn, chunk_id)
            if value is None:
                raise _EarlyReturn(
                    error_response("NOT_FOUND", "Chunk not found", status_code=404)
                )

            tags = _get_tags(value)
            target = next((t for t in tags if t.get("tag_id") == tag_id), None)
            if target is None:
                raise _EarlyReturn(
                    error_response("NOT_FOUND", f"Tag {tag_id} not found", status_code=404)
                )

            # Check for duplicate tag name on same chunk
            if any(t.get("name") == body.name and t.get("tag_id") != tag_id for t in tags):
                raise _EarlyReturn(
                    error_response(
                        "DUPLICATE_TAG_NAME",
                        f"Tag name '{body.name}' already exists on this chunk",
                        status_code=400,
                    )
                )

            old_name = target["name"]
            target["name"] = body.name
            target.setdefault("history", []).append(
                _make_history_entry("edited", user_id, old_name=old_name)
            )

            _write_tags(conn, chunk_id, tags)
            upsert_chunk_field(
                chunk_id,
                tags=tags,
                verified_tag_names=project_verified_tag_names(tags, chunk_id),
            )

        resp = success_response({
            "chunk_id": chunk_id,
            "tag_id": tag_id,
            "old_name": old_name,
            "new_name": body.name,
        })
        record_outcome(tag_op_total, resp.status_code, operation=_op, tenant_id=tenant_id)
        tag_endpoint_duration_seconds.labels(operation=_op).observe(time.perf_counter() - _t0)
        return resp

    except _EarlyReturn as e:
        record_outcome(tag_op_total, e.response.status_code, operation=_op, tenant_id=tenant_id)
        tag_endpoint_duration_seconds.labels(operation=_op).observe(time.perf_counter() - _t0)
        return e.response
    except Exception as e:
        log.exception("Tag rename failed: %s", e)
        resp = error_response("INTERNAL_ERROR", "Tag rename failed", status_code=500)
        record_outcome(tag_op_total, resp.status_code, operation=_op, tenant_id=tenant_id)
        tag_endpoint_duration_seconds.labels(operation=_op).observe(time.perf_counter() - _t0)
        return resp


@tags_app.delete("/{chunk_id}/tags/{tag_id}", status_code=204)
def delete_tag(chunk_id: str, tag_id: str, request: Request):
    """Hard-delete a tag from a chunk."""
    ctx = request.state.hmac_ctx
    tenant_id = ctx["tenant_id"]
    _t0 = time.perf_counter()
    _op = "delete"

    try:
        with rls_connection(tenant_id) as conn:
            value = _read_chunk_for_update(conn, chunk_id)
            if value is None:
                raise _EarlyReturn(
                    error_response("NOT_FOUND", "Chunk not found", status_code=404)
                )

            tags = _get_tags(value)
            updated_tags = [t for t in tags if t.get("tag_id") != tag_id]
            if len(updated_tags) == len(tags):
                raise _EarlyReturn(
                    error_response("NOT_FOUND", f"Tag {tag_id} not found", status_code=404)
                )

            _write_tags(conn, chunk_id, updated_tags)
            upsert_chunk_field(
                chunk_id,
                tags=updated_tags,
                verified_tag_names=project_verified_tag_names(updated_tags, chunk_id),
            )

        record_outcome(tag_op_total, 204, operation=_op, tenant_id=tenant_id)
        tag_endpoint_duration_seconds.labels(operation=_op).observe(time.perf_counter() - _t0)
        return Response(status_code=204)

    except _EarlyReturn as e:
        record_outcome(tag_op_total, e.response.status_code, operation=_op, tenant_id=tenant_id)
        tag_endpoint_duration_seconds.labels(operation=_op).observe(time.perf_counter() - _t0)
        return e.response
    except Exception as e:
        log.exception("Tag deletion failed: %s", e)
        resp = error_response("INTERNAL_ERROR", "Tag deletion failed", status_code=500)
        record_outcome(tag_op_total, resp.status_code, operation=_op, tenant_id=tenant_id)
        tag_endpoint_duration_seconds.labels(operation=_op).observe(time.perf_counter() - _t0)
        return resp
