import json

# Query Parser Prompt

QUERY_SYSTEM_PROMPT = """
You are a query-understanding component for an image search engine.

Correct obvious spelling errors in the user's Arabic or English query,
then return ONLY one valid JSON object with this exact structure:

{
  "language": "",
  "corrected_query": "",
  "semantic_query": "",
  "object_terms": [],
  "text_terms": [],
  "attributes": [],
  "relations": [],
  "search_mode": ""
}

Rules:

1. language:
   - Use "ar" for Arabic queries.
   - Use "en" for English queries.

2. corrected_query:
   - Correct spelling and grammar.
   - Keep it in the same language as the original query.
   - Do not translate it.
   - Preserve the user's dialect when possible.

3. semantic_query:
   - Write a clear English description representing the user's meaning.
   - Do not add details that the user did not mention.

4. object_terms:
   - Important visible objects.
   - Use English singular words.

5. text_terms:
   - Exact text that the user wants to find inside images.
   - Preserve capitalization and wording.

6. attributes:
   - Colors and visual characteristics in English.

7. relations:
   - Relationships between objects in English.

8. search_mode must be one of:
   - "semantic"
   - "object"
   - "text"
   - "hybrid"

9. Use "hybrid" when the query includes objects with attributes,
   relations, or more than one search method.

10. Return JSON only.
Do not use markdown and do not include explanations.
"""


# Run Query Parser using Qwen

def run_qwen_text(user_query, vlm_model, vlm_processor):
    """
    Convert the user's natural-language query
    into a structured search query using Qwen.
    """

    messages = [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": QUERY_SYSTEM_PROMPT
                }
            ]
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": user_query
                }
            ]
        }
    ]

    inputs = vlm_processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt"
    )

    inputs = inputs.to(vlm_model.device)

    generated_ids = vlm_model.generate(
        **inputs,
        max_new_tokens=300,
        do_sample=False,
        pad_token_id=vlm_processor.tokenizer.eos_token_id
    )

    generated_ids = generated_ids[
        :,
        inputs["input_ids"].shape[1]:
    ]

    output_text = vlm_processor.batch_decode(
        generated_ids,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False
    )[0]

    return output_text


# Clean Qwen JSON Output

def prepare_json_text(raw_text):
    """
    Remove markdown formatting and extract
    the JSON object from Qwen's response.
    """

    cleaned_text = raw_text.strip()

    cleaned_text = cleaned_text.replace(
        "```json",
        ""
    )

    cleaned_text = cleaned_text.replace(
        "```",
        ""
    )

    cleaned_text = cleaned_text.strip()

    start_index = cleaned_text.find("{")
    end_index = cleaned_text.rfind("}")

    if start_index == -1 or end_index == -1:
        raise ValueError(
            "Qwen output does not contain a valid JSON object."
        )

    return cleaned_text[
        start_index:end_index + 1
    ]


# Parse User Query

def parse_search_query(
    user_query,
    vlm_model,
    vlm_processor
):
    """
    Convert a user's search query into structured fields.
    """

    raw_output = run_qwen_text(
        user_query,
        vlm_model,
        vlm_processor
    )

    cleaned_output = prepare_json_text(
        raw_output
    )

    try:
        parsed_query = json.loads(
            cleaned_output
        )

    except json.JSONDecodeError as error:

        print("Raw Qwen output:")
        print(raw_output)

        raise ValueError(
            f"Qwen returned invalid JSON: {error}"
        )

    required_fields = [
        "language",
        "corrected_query",
        "semantic_query",
        "object_terms",
        "text_terms",
        "attributes",
        "relations",
        "search_mode"
    ]

    for field in required_fields:

        if field not in parsed_query:

            raise ValueError(
                f"Missing query field: {field}"
            )

    return parsed_query