data_studio/api.py

from __future__ import annotations

from typing import Any, Literal

from fastapi import Depends, FastAPI, Header, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from data_studio.config import ConfigError, StudioSettings
from data_studio.auth import create_or_update_user, generate_temporary_password, get_user
from data_studio.db import create_client, ensure_indexes, get_database
from data_studio.exports import export_clean_dataset
from data_studio.benchmarks import assert_no_exact_benchmark_matches, register_protected_benchmark
from data_studio.governance import SourceRegistryV1, SourceState
from data_studio.governed_importer import import_governed_record
from data_studio.governed_sources import get_registered_source, register_source, transition_source
from data_studio.importer import import_source_samples
from data_studio.rubrics import RubricV1, register_rubric
from data_studio.review import (
    ReviewConflict,
    ReviewValidationError,
    apply_review_action,
    claim_next_sample_with_enrichment,
    queue_overview,
    release_claim,
)
from data_studio.security import create_session, delete_session, get_session_user, has_permission
from data_studio.sources import list_sources, seed_default_sources
from data_studio.utils import utc_now_iso


settings = StudioSettings.from_env()
client = create_client(settings)
db = get_database(client, settings)
ensure_indexes(db)

app = FastAPI(title="Rachana Data Studio API", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:3000", "http://127.0.0.1:3000"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.on_event("startup")
def startup_seed_sources() -> None:
    ensure_indexes(db)
    seed_default_sources(db)


class LoginRequest(BaseModel):
    username: str
    password: str


class ReviewRequest(BaseModel):
    action: Literal["accept", "edit", "reject", "skip"]
    expected_version: int
    cleaned_payload: dict[str, Any]
    tags: list[str] = Field(default_factory=list)
    quality_tags: list[str] = Field(default_factory=list)
    task_tags: list[str] = Field(default_factory=list)
    notes: str = ""
    confidence: int = Field(default=4, ge=1, le=5)
    reason: str = ""


class ImportRequest(BaseModel):
    dataset_key: str
    max_samples: int = Field(default=100, ge=1, le=50_000)


class ExportRequest(BaseModel):
    version: str = Field(min_length=1, max_length=80, pattern=r"^[A-Za-z0-9._-]+$")
    repo_id: str | None = None


class RegisterSourceRequest(BaseModel):
    source: SourceRegistryV1


class TransitionSourceRequest(BaseModel):
    target: SourceState
    reason: str = Field(min_length=1)


class RegisterBenchmarkRequest(BaseModel):
    benchmark_id: str = Field(min_length=1)
    exact_revision: str = Field(min_length=1)
    task: str = Field(min_length=1)
    exact_payloads: list[dict[str, Any]]


class RegisterRubricRequest(BaseModel):
    rubric: RubricV1


class GovernedRecordRequest(BaseModel):
    source_id: str = Field(min_length=1)
    source_revision: str = Field(min_length=1)
    source_native_id: str = Field(min_length=1)
    raw_payload: dict[str, Any]
    normalized_payload: dict[str, Any]


class CreateUserRequest(BaseModel):
    username: str = Field(min_length=1)
    password: str = Field(min_length=8)
    display_name: str | None = None
    role: Literal["reviewer", "manager", "admin"] = "reviewer"
    update_existing: bool = False
    inactive: bool = False


class ResetPasswordRequest(BaseModel):
    temporary_password: str | None = Field(default=None, min_length=8)


class ResetStudioRequest(BaseModel):
    confirm: Literal["RESET"]


class ResetReviewPoolRequest(BaseModel):
    confirm: Literal["RESET_REVIEW_POOL"]


def bearer_token(authorization: str = Header(default="")) -> str:
    scheme, _, token = authorization.partition(" ")
    if scheme.lower() != "bearer" or not token:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Authentication required.")
    return token


def current_user(token: str = Depends(bearer_token)) -> dict[str, Any]:
    user = get_session_user(db, token)
    if user is None:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Session expired.")
    return user


def require_permission(permission: str):
    def dependency(user: dict[str, Any] = Depends(current_user)) -> dict[str, Any]:
        if not has_permission(user, permission):
            raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=f"Missing permission: {permission}")
        return user

    return dependency


@app.get("/api/health")
def health() -> dict[str, Any]:
    return {"status": "ok", **settings.deployment_metadata()}


@app.post("/api/auth/login")
def login(payload: LoginRequest) -> dict[str, Any]:
    result = create_session(db, payload.username, payload.password)
    if result is None:
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid username or password.")
    token, user = result
    return {"token": token, "user": user}


@app.post("/api/auth/logout")
def logout(token: str = Depends(bearer_token)) -> dict[str, bool]:
    delete_session(db, token)
    return {"ok": True}


@app.get("/api/me")
def me(user: dict[str, Any] = Depends(current_user)) -> dict[str, Any]:
    return user


@app.get("/api/overview")
def overview(user: dict[str, Any] = Depends(current_user)) -> dict[str, Any]:
    return {**queue_overview(db, settings), "source_configs": list_sources(db), "user": user}


@app.post("/api/review/claim/{queue_name}")
def claim(
    queue_name: str,
    dataset_key: str | None = None,
    user: dict[str, Any] = Depends(require_permission("review")),
) -> dict[str, Any]:
    return {"sample": claim_next_sample_with_enrichment(db, settings, queue_name, user["username"], dataset_key)}


@app.post("/api/review/{sample_id}/release")
def release(sample_id: str, user: dict[str, Any] = Depends(require_permission("review"))) -> dict[str, bool]:
    release_claim(db, sample_id, user["username"])
    return {"ok": True}


@app.post("/api/review/{sample_id}")
def review(
    sample_id: str,
    payload: ReviewRequest,
    user: dict[str, Any] = Depends(require_permission("review")),
) -> dict[str, Any]:
    try:
        resolved_action = apply_review_action(
            db=db,
            settings=settings,
            sample_id=sample_id,
            reviewer=user["username"],
            expected_version=payload.expected_version,
            action=payload.action,
            cleaned_payload=payload.cleaned_payload,
            tags=payload.tags,
            quality_tags=payload.quality_tags,
            task_tags=payload.task_tags,
            notes=payload.notes,
            confidence=payload.confidence,
            reason=payload.reason,
        )
    except ReviewConflict as exc:
        raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(exc)) from exc
    except ReviewValidationError as exc:
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)) from exc
    return resolved_action


@app.post("/api/admin/sources/seed")
def seed_sources(user: dict[str, Any] = Depends(require_permission("manage_sources"))) -> dict[str, int]:
    return {"inserted": seed_default_sources(db)}


@app.post("/api/admin/reset")
def reset_studio(
    payload: ResetStudioRequest,
    user: dict[str, Any] = Depends(require_permission("manage_users")),
) -> dict[str, Any]:
    collections = [
        "samples",
        "review_events",
        "import_jobs",
        "import_checkpoints",
        "exports",
        "dataset_versions",
        "sessions",
        "source_configs",
    ]
    cleared = {}
    for name in collections:
        result = db[name].delete_many({})
        cleared[name] = result.deleted_count
    inserted_sources = seed_default_sources(db)
    return {"cleared": cleared, "seeded_sources": inserted_sources, "confirmed": payload.confirm}


@app.post("/api/admin/review/reset")
def reset_review_pool(
    payload: ResetReviewPoolRequest,
    user: dict[str, Any] = Depends(require_permission("manage_users")),
) -> dict[str, Any]:
    reviewed_statuses = ["accepted", "edited", "rejected", "skipped"]
    now = utc_now_iso()
    reset_counts = {"samples": 0, "review_events": 0}

    reviewed_samples = list(db["samples"].find({"status": {"$in": reviewed_statuses}}))
    for sample in reviewed_samples:
        sample_type = sample["sample_type"]
        review_reset = {
            "current_text": None,
            "current_pair": None,
            "current_transliteration": None,
            "tags": [],
            "quality_tags": [],
            "task_tags": [],
            "notes": "",
            "last_reviewed_by": None,
            "last_reviewed_at": None,
            "edit_count": 0,
            "token_stats": None,
        }
        if sample_type == "document":
            review_reset["current_text"] = sample["payload"].get("text", "")
        elif sample_type == "translation_pair":
            review_reset["current_pair"] = {
                "source_text": sample["payload"].get("source_text", ""),
                "target_text": sample["payload"].get("target_text", ""),
            }
        elif sample_type == "transliteration_pair":
            review_reset["current_transliteration"] = {
                "native_text": sample["payload"].get("native_text", ""),
                "latin_text": sample["payload"].get("latin_text", ""),
            }
        db["samples"].update_one(
            {"sample_id": sample["sample_id"]},
            {
                "$set": {
                    "status": "pending",
                    "review": review_reset,
                    "updated_at": now,
                    "metrics.clean_char_count": None,
                    "metrics.clean_token_count": None,
                    "metrics.tokenizer_version": None,
                    "export.eligible": False,
                    "governance.review_accepted": False,
                    "governance.training_export_eligible": False,
                    "live_sync.repo_id": None,
                    "live_sync.path_in_repo": None,
                    "live_sync.status": "pending",
                    "live_sync.last_synced_at": None,
                    "live_sync.last_error": None,
                },
                "$unset": {"lease": ""},
                "$inc": {"version": 1},
            },
        )
        reset_counts["samples"] += 1

    deleted_events = db["review_events"].delete_many({"sample_id": {"$in": [sample["sample_id"] for sample in reviewed_samples]}})
    reset_counts["review_events"] = deleted_events.deleted_count

    return {
        "confirmed": payload.confirm,
        "reset_counts": reset_counts,
        "metadata": now,
    }


@app.get("/api/admin/users")
def list_users(user: dict[str, Any] = Depends(require_permission("manage_users"))) -> dict[str, Any]:
    users = []
    for row in db["users"].find().sort("username", 1):
        users.append(
            {
                "username": row["username"],
                "display_name": row.get("display_name") or row["username"],
                "role": row.get("role", "reviewer"),
                "is_active": bool(row.get("is_active", True)),
                "last_login_at": row.get("last_login_at"),
            }
        )
    return {"users": users}


@app.post("/api/admin/users")
def create_user(
    payload: CreateUserRequest,
    user: dict[str, Any] = Depends(require_permission("manage_users")),
) -> dict[str, Any]:
    try:
        created = create_or_update_user(
            db=db,
            username=payload.username,
            password=payload.password,
            display_name=payload.display_name,
            role=payload.role,
            is_active=not payload.inactive,
            update_existing=payload.update_existing,
        )
    except ValueError as exc:
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)) from exc
    return {"user": created}


@app.post("/api/admin/users/{username}/reset-password")
def reset_user_password(
    username: str,
    payload: ResetPasswordRequest,
    user: dict[str, Any] = Depends(require_permission("manage_users")),
) -> dict[str, Any]:
    existing = get_user(db, username)
    if existing is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found.")
    temporary_password = payload.temporary_password or generate_temporary_password()
    try:
        created = create_or_update_user(
            db=db,
            username=username,
            password=temporary_password,
            display_name=existing.get("display_name"),
            role=existing.get("role", "reviewer"),
            is_active=bool(existing.get("is_active", True)),
            update_existing=True,
        )
    except ValueError as exc:
        raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)) from exc
    return {"user": created, "temporary_password": temporary_password}


@app.post("/api/admin/import")
def run_import(
    payload: ImportRequest,
    user: dict[str, Any] = Depends(require_permission("import")),
) -> dict[str, Any]:
    try:
        settings.require_legacy_import_export_enabled()
    except ConfigError as exc:
        raise HTTPException(status_code=status.HTTP_423_LOCKED, detail=str(exc)) from exc
    return import_source_samples(db, settings, payload.dataset_key, payload.max_samples)


@app.post("/api/admin/export")
def run_export(
    payload: ExportRequest,
    user: dict[str, Any] = Depends(require_permission("export")),
) -> dict[str, Any]:
    try:
        settings.require_legacy_import_export_enabled()
    except ConfigError as exc:
        raise HTTPException(status_code=status.HTTP_423_LOCKED, detail=str(exc)) from exc
    return export_clean_dataset(db, settings, payload.version, user["username"], payload.repo_id)


@app.post("/api/governance/sources")
def create_governed_source(
    payload: RegisterSourceRequest,
    user: dict[str, Any] = Depends(require_permission("manage_sources")),
) -> dict[str, str]:
    register_source(db, payload.source, user["username"])
    return {"source_id": payload.source.source_id, "revision": payload.source.hf_revision}


@app.post("/api/governance/sources/{source_id}/{revision}/transition")
def change_governed_source_state(
    source_id: str,
    revision: str,
    payload: TransitionSourceRequest,
    user: dict[str, Any] = Depends(require_permission("approve_sources")),
) -> dict[str, str]:
    transition_source(db, source_id, revision, payload.target, user["username"], payload.reason)
    return {"source_id": source_id, "revision": revision, "state": payload.target.value}


@app.post("/api/governance/benchmarks")
def create_protected_benchmark(
    payload: RegisterBenchmarkRequest,
    user: dict[str, Any] = Depends(require_permission("manage_benchmarks")),
) -> dict[str, str]:
    register_protected_benchmark(
        db,
        payload.benchmark_id,
        payload.exact_revision,
        payload.task,
        payload.exact_payloads,
        user["username"],
    )
    return {"benchmark_id": payload.benchmark_id, "revision": payload.exact_revision}


@app.post("/api/governance/rubrics")
def create_rubric(
    payload: RegisterRubricRequest,
    user: dict[str, Any] = Depends(require_permission("manage_rubrics")),
) -> dict[str, str]:
    register_rubric(db, payload.rubric, user["username"])
    return {"rubric_id": payload.rubric.rubric_id, "rubric_version": payload.rubric.rubric_version}


@app.post("/api/governance/records")
def create_governed_record(
    payload: GovernedRecordRequest,
    user: dict[str, Any] = Depends(require_permission("governed_import")),
) -> dict[str, Any]:
    source = get_registered_source(db, payload.source_id, payload.source_revision)
    if source is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Registered source revision not found.")
    assert_no_exact_benchmark_matches(db, [payload.raw_payload, payload.normalized_payload])
    record_id, inserted = import_governed_record(
        db,
        source,
        payload.source_native_id,
        payload.raw_payload,
        payload.normalized_payload,
    )
    return {"record_id": record_id, "inserted": inserted}