app/api/v1/json_extract.py

from __future__ import annotations

import asyncio
import time
from typing import Annotated, Any, Dict, List, Optional

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field

from app.config import get_settings
from app.core.logger import get_logger
from app.core.thread_pool import run_in_executor
from app.services.gliner_service import gliner_service
from app.services.json_service import extract_json
from app.services.monetary_field_service import MonetaryFieldStatus, apply_monetary_fields

logger = get_logger(__name__)

router = APIRouter()

MAX_CONTENT_LENGTH = 10_000_000


class ExtractJsonRequest(BaseModel):
    content: str = Field(
        ...,
        description="Dirty string content potentially containing JSON wrapped in markdown, conversational text, etc.",
        min_length=1,
    )
    limit: Optional[int] = Field(
        default=None,
        ge=1,
        le=100,
        description="Maximum number of JSON objects to extract. Omit for all.",
    )
    mode: str = Field(
        default="all",
        pattern=r"^(first|all)$",
        description="'first' returns only the first JSON object; 'all' returns all extracted objects.",
    )


class ExtractJsonResponse(BaseModel):
    success: bool
    time_ms: float
    data: Any = None
    count: int = 0
    error_message: Optional[str] = None


@router.post(
    "/json/extract",
    response_model=ExtractJsonResponse,
    summary="Extract JSON from dirty/markdown content",
    description=(
        "Accepts string content that may contain JSON embedded in markdown code fences "
        "(```json), XML-style <json> tags, or mixed with conversational text. "
        "Returns cleaned, parsed JSON objects. Handles malformed JSON via a repair pipeline "
        "that fixes trailing commas, unquoted keys, single-quote strings, JS comments, etc."
    ),
)
async def extract_json_endpoint(
    body: ExtractJsonRequest,
) -> ExtractJsonResponse:
    start = time.perf_counter()

    content_length = len(body.content)
    if content_length > MAX_CONTENT_LENGTH:
        elapsed = round((time.perf_counter() - start) * 1000, 3)
        raise HTTPException(
            status_code=413,
            detail=ExtractJsonResponse(
                success=False,
                time_ms=elapsed,
                data=None,
                count=0,
                error_message=f"Content exceeds maximum length of {MAX_CONTENT_LENGTH:,} characters.",
            ).model_dump(),
        )

    effective_limit = 1 if body.mode == "first" else body.limit

    result = await run_in_executor(extract_json, body.content, limit=effective_limit)

    elapsed = round((time.perf_counter() - start) * 1000, 3)

    if not result.success:
        logger.warning(
            "JSON extraction returned no results",
            extra={
                "input_length": content_length,
                "mode": body.mode,
                "time_ms": elapsed,
            },
        )
        raise HTTPException(
            status_code=422,
            detail=ExtractJsonResponse(
                success=False,
                time_ms=elapsed,
                data=None,
                count=0,
                error_message=result.error_message or "No JSON content could be extracted from the provided input.",
            ).model_dump(),
        )

    response_data = result.data[0] if body.mode == "first" else result.data

    logger.info(
        "JSON extraction successful",
        extra={
            "count": result.total_extracted,
            "method": result.extraction_method,
            "input_length": content_length,
            "time_ms": elapsed,
        },
    )

    return ExtractJsonResponse(
        success=True,
        time_ms=elapsed,
        data=response_data,
        count=result.total_extracted,
        error_message=None,
    )


class NoAiExtractRequest(BaseModel):
    content: str = Field(
        ...,
        description=(
            "Raw text to extract from (invoice text, OCR output, emails, etc.). "
            "No external AI/LLM API is called."
        ),
        min_length=1,
    )
    mode: str = Field(
        default="json",
        pattern=r"^(json|entities)$",
        description=(
            "'json' extracts structured fields using a GLiNER2 `structure` schema; "
            "'entities' extracts zero-shot entities using a `labels` list."
        ),
    )
    structure: Optional[Dict[str, Any]] = Field(
        default=None,
        description=(
            "Required for mode='json'. GLiNER2 structure schema mapping a parent "
            "key to field specs, e.g. "
            '{"invoice": ["number::str::Invoice number", "total::str::Total amount"]}. '
            "Field spec: name::dtype::choices::description."
        ),
    )
    labels: Optional[List[str]] = Field(
        default=None,
        description="Required for mode='entities'. Entity types to detect, e.g. ['person', 'company', 'location'].",
    )
    threshold: float = Field(
        default=0.5,
        ge=0.0,
        le=1.0,
        description="Confidence threshold (0.0-1.0). Lower includes more candidates.",
    )
    monetary_fields: Optional["MonetaryFieldsConfig"] = Field(
        default=None,
        description=(
            "OPTIONAL. Declare which extracted fields hold monetary values so they "
            "are price-parsed in place: `field_names` lists the monetary fields "
            "inside each object (e.g. ['total']) and `object_name` selects the "
            "container key in the extracted `data` -- it is inferred from the "
            "`structure` parent key when omitted (mode='json') and validated "
            "against the structure when provided. Found values are replaced with "
            "their parsed numeric amount; missing/blank/unparseable values are "
            "reported per field. Omit to leave the extracted data untouched."
        ),
    )


class MonetaryFieldsConfig(BaseModel):
    """Declares which extracted fields are monetary so they are price-parsed."""

    object_name: Optional[str] = Field(
        default=None,
        description=(
            "Container key in the extracted `data` whose value holds the objects to "
            "process, e.g. 'invoice'. OPTIONAL for mode='json': inferred from the "
            "single parent key of `structure`. When provided it must match a "
            "structure object name. Required for mode='entities'."
        ),
    )
    field_names: List[str] = Field(
        ...,
        min_length=1,
        description="Monetary field names inside each object, e.g. ['total', 'cgst'].",
    )


class NoAiExtractResponse(BaseModel):
    success: bool
    time_ms: float
    mode: str
    data: Any = None
    count: int = 0
    error_message: Optional[str] = None
    monetary_fields: List[MonetaryFieldStatus] = Field(
        default_factory=list,
        description=(
            "Per-field outcome of the optional `monetary_fields` config: status "
            "is 'parsed', 'not_found', 'not_parsable' or 'skipped', with the "
            "raw `value`, the `parsed_value` and a human-readable error. Empty "
            "when no config was sent."
        ),
    )


def _count_extracted(result: Any) -> int:
    if not isinstance(result, dict):
        return 0
    total = 0
    for parent, items in result.items():
        if isinstance(items, list):
            total += len(items)
        elif isinstance(items, dict):
            total += 1
    return total


@router.post(
    "/json/feature-extract",
    response_model=List[NoAiExtractResponse],
    summary="Batch-extract structured JSON / entities with a local on-device model (no external AI)",
    description=(
        "Send up to 5 requests as a JSON array (1-5 items). All items are "
        "processed concurrently on the shared thread pool, bounded by the local "
        "model's concurrency limit. No external AI/LLM API is contacted -- ideal "
        "for private or invoice/OCR data. mode='json' uses a structure schema to "
        "pull named fields; mode='entities' detects a flat list of entity types. "
        "Each item reports its own success/error. "
        "Returns HTTP 503 if the model failed to load or is disabled. "
        "OPTIONAL per-item `monetary_fields` ({\"object_name\", \"field_names\"}) "
        "price-parses the declared fields in place inside the extracted data."
    ),
)
async def no_ai_extract_endpoint(
    body: Annotated[
        List[NoAiExtractRequest],
        Field(
            min_length=1,
            max_length=5,
            description="Array of up to 5 extraction requests, processed concurrently.",
        ),
    ],
) -> List[NoAiExtractResponse]:
    settings = get_settings()
    total_start = time.perf_counter()

    if not settings.gliner_enabled:
        raise HTTPException(
            status_code=503,
            detail=NoAiExtractResponse(
                success=False,
                time_ms=0.0,
                mode="",
                data=None,
                count=0,
                error_message="The local extraction service is disabled.",
            ).model_dump(),
        )

    if not gliner_service.is_loaded():
        try:
            await run_in_executor(gliner_service.load_model)
        except Exception:
            logger.exception("Lazy GLiNER2 model load failed on request")
            raise HTTPException(
                status_code=503,
                detail=NoAiExtractResponse(
                    success=False,
                    time_ms=round((time.perf_counter() - total_start) * 1000, 3),
                    mode="",
                    data=None,
                    count=0,
                    error_message="The local extraction model is unavailable. Please try again later.",
                ).model_dump(),
            )

    async def _process_one(index: int, item: NoAiExtractRequest) -> NoAiExtractResponse:
        start = time.perf_counter()

        if item.mode == "json" and not item.structure:
            return NoAiExtractResponse(
                success=False,
                time_ms=round((time.perf_counter() - start) * 1000, 3),
                mode=item.mode,
                data=None,
                count=0,
                error_message="mode='json' requires a non-empty 'structure' schema.",
            )
        if item.mode == "entities" and not item.labels:
            return NoAiExtractResponse(
                success=False,
                time_ms=round((time.perf_counter() - start) * 1000, 3),
                mode=item.mode,
                data=None,
                count=0,
                error_message="mode='entities' requires a non-empty 'labels' list.",
            )
        if len(item.content) > gliner_service.max_content_length:
            return NoAiExtractResponse(
                success=False,
                time_ms=round((time.perf_counter() - start) * 1000, 3),
                mode=item.mode,
                data=None,
                count=0,
                error_message=(
                    f"Content exceeds maximum length of {gliner_service.max_content_length:,} characters."
                ),
            )

        # Resolve the monetary-fields target object name. It is optional for
        # mode='json' (inferred from the single `structure` parent key) and
        # validated against the structure when provided.
        monetary_config = item.monetary_fields
        object_name: Optional[str] = None
        if monetary_config is not None:
            object_name = monetary_config.object_name
            if item.mode == "json" and isinstance(item.structure, dict):
                structure_names = [k for k in item.structure if isinstance(k, str)]
                if object_name is None:
                    if len(structure_names) == 1:
                        object_name = structure_names[0]
                    else:
                        return NoAiExtractResponse(
                            success=False,
                            time_ms=round((time.perf_counter() - start) * 1000, 3),
                            mode=item.mode,
                            data=None,
                            count=0,
                            error_message=(
                                "monetary_fields.object_name is required when the structure "
                                "has multiple object names: " + ", ".join(structure_names)
                            ),
                        )
                elif object_name not in structure_names:
                    return NoAiExtractResponse(
                        success=False,
                        time_ms=round((time.perf_counter() - start) * 1000, 3),
                        mode=item.mode,
                        data=None,
                        count=0,
                        error_message=(
                            f"monetary_fields.object_name '{object_name}' does not match "
                            "the structure object name(s): " + ", ".join(structure_names)
                        ),
                    )
            elif item.mode == "entities" and object_name is None:
                return NoAiExtractResponse(
                    success=False,
                    time_ms=round((time.perf_counter() - start) * 1000, 3),
                    mode=item.mode,
                    data=None,
                    count=0,
                    error_message=(
                        "monetary_fields.object_name is required for mode='entities' "
                        "(no structure to infer it from)"
                    ),
                )

        try:
            if item.mode == "json":
                result = await run_in_executor(
                    gliner_service.extract_json, item.content, item.structure, item.threshold
                )
            else:
                result = await run_in_executor(
                    gliner_service.extract_entities, item.content, item.labels, item.threshold
                )
        except Exception:
            logger.exception("GLiNER2 inference failed for item %s", index)
            return NoAiExtractResponse(
                success=False,
                time_ms=round((time.perf_counter() - start) * 1000, 3),
                mode=item.mode,
                data=None,
                count=0,
                error_message="Extraction failed. Please try again later.",
            )

        elapsed = round((time.perf_counter() - start) * 1000, 3)

        # OPTIONAL post-processing: price-parse the declared monetary fields
        # in place inside the extracted data (e.g. "1250.75" -> 1250.75) and
        # report the per-field outcome so callers can understand any failures.
        monetary_report: List[MonetaryFieldStatus] = []
        parsed_monetary = 0
        if item.monetary_fields is not None:
            monetary_report = apply_monetary_fields(
                result,
                object_name or "",
                item.monetary_fields.field_names,
            )
            parsed_monetary = sum(1 for s in monetary_report if s.status == "parsed")

        logger.info(
            "Feature-extract item extracted",
            extra={
                "index": index,
                "mode": item.mode,
                "count": _count_extracted(result),
                "monetary_fields_parsed": parsed_monetary,
                "time_ms": elapsed,
            },
        )
        return NoAiExtractResponse(
            success=True,
            time_ms=elapsed,
            mode=item.mode,
            data=result,
            count=_count_extracted(result),
            error_message=None,
            monetary_fields=monetary_report,
        )

    results = await asyncio.gather(
        *[_process_one(index, item) for index, item in enumerate(body)]
    )

    logger.info(
        "Feature-extract batch processed",
        extra={
            "items": len(body),
            "time_ms": round((time.perf_counter() - total_start) * 1000, 3),
        },
    )
    return list(results)