"""SHARED CONTRACT — the prompt + output schema the model is trained and served on.

This file is the single source of truth for:
  - FIELD_ORDER          : canonical key order of the extracted JSON
  - build_prompt()       : the EXACT prompt string used at train AND inference time
  - format_completion()  : training target serialization (also used by eval)
  - extract_json()       : brace-matched JSON extraction from raw model text
  - HotelExtraction      : Pydantic schema for validation
  - validate()           : parse+validate -> clean dict (or None)

It is intentionally DEPENDENCY-LIGHT (only json + pydantic) so both the heavy
training env and the lean serving env can import it. A COPY of this file lives in
BOTH training/ and deployment/; tests/test_contract_parity.py asserts they are
byte-identical, so the prompt can never silently drift between train and serve.

DO NOT EDIT one copy without the other. Edit the source, re-sync, re-run the test.
"""

from __future__ import annotations

import json
from typing import Any, Literal, Optional

from pydantic import BaseModel, ConfigDict, Field, ValidationError

DEFAULT_BASE_MODEL = "Qwen/Qwen3-1.7B"

# `filters` is emitted as a comma-separated string of NATURAL PHRASES
# (e.g. "swimming pool, pet friendly"), NOT production codes. A downstream matcher
# resolves phrases -> codes (FL_HF_29, ...), keeping the model taxonomy-agnostic.
FIELD_ORDER = [
    "destination", "locality", "hotelName", "checkinDate", "checkoutDate",
    "adultCount", "roomCount", "childCount", "childAges", "infantCount",
    "sortCriteria", "minPrice", "maxPrice", "filters", "deepSearch", "isNearMe", "resetAction",
]


def build_prompt(query: str, today: str) -> str:
    """The exact user-message content used for BOTH training and inference."""
    return (
        "Extract hotel-search entities.\n"
        "Return strict JSON only.\n"
        "Schema:\n"
        "destination, locality, hotelName, checkinDate, checkoutDate, "
        "adultCount, roomCount, childCount, childAges, infantCount, "
        "sortCriteria, minPrice, maxPrice, filters, deepSearch, isNearMe, resetAction.\n\n"
        "Rules:\n"
        "- Dates are DDMMYYYY.\n"
        '- deepSearch and isNearMe must be "true" or "false".\n'
        "- filters is a comma-separated list of amenity/type phrases (e.g. "
        '"swimming pool, pet friendly"); prefix removals with "no ".\n'
        "- Omit unknown optional fields.\n\n"
        f"today={today}\n"
        f"query={query}"
    )


def format_completion(expected: dict[str, Any]) -> str:
    """Stable-key-order minified JSON — the training target / gold serialization."""
    ordered = {f: expected[f] for f in FIELD_ORDER if f in expected}
    for k in expected:
        if k not in ordered:
            ordered[k] = expected[k]
    return json.dumps(ordered, ensure_ascii=False, separators=(",", ":"))


def extract_json(text: str) -> dict | None:
    """Pull the first balanced {...} object out of raw model text."""
    start = text.find("{")
    if start < 0:
        return None
    depth = 0
    in_str = esc = False
    for i in range(start, len(text)):
        c = text[i]
        if esc:
            esc = False
            continue
        if c == "\\":
            esc = True
            continue
        if c == '"':
            in_str = not in_str
            continue
        if in_str:
            continue
        if c == "{":
            depth += 1
        elif c == "}":
            depth -= 1
            if depth == 0:
                try:
                    return json.loads(text[start:i + 1])
                except json.JSONDecodeError:
                    return None
    return None


class HotelExtraction(BaseModel):
    """Output contract. `filters` is a comma-separated PHRASE string (e.g.
    "swimming pool, pet friendly"); a downstream matcher resolves it to codes."""
    model_config = ConfigDict(extra="forbid")
    deepSearch: Literal["true", "false"]
    isNearMe: Literal["true", "false"]
    destination: Optional[str] = Field(default=None, min_length=1)
    locality: Optional[str] = Field(default=None, min_length=1)
    hotelName: Optional[str] = Field(default=None, min_length=1)
    checkinDate: Optional[str] = Field(default=None, pattern=r"^\d{8}$")
    checkoutDate: Optional[str] = Field(default=None, pattern=r"^\d{8}$")
    adultCount: Optional[int] = Field(default=None, ge=0, le=20)
    roomCount: Optional[int] = Field(default=None, ge=0, le=20)
    childCount: Optional[int] = Field(default=None, ge=0, le=20)
    childAges: Optional[list[int]] = None
    infantCount: Optional[int] = Field(default=None, ge=0, le=20)
    sortCriteria: Optional[Literal["SC_P_LH", "SC_P_HL", "SC_UR", "SC_P", "SC_DIST"]] = None
    minPrice: Optional[int] = Field(default=None, ge=0, le=10_000_000)
    maxPrice: Optional[int] = Field(default=None, ge=0, le=10_000_000)
    filters: Optional[str] = Field(default=None, min_length=1)
    resetAction: Optional[Literal["filters", "guests", "dates", "all"]] = None


def validate(d: dict) -> dict | None:
    """Validate a parsed dict against the schema; return clean dict or None."""
    try:
        return HotelExtraction.model_validate(d).model_dump(exclude_none=True)
    except ValidationError:
        return None