import json
import torch

from vision_model import get_model, get_processor

EXTRACTION_PROMPT = """
Return ONLY valid JSON.

Schema:
{
  "category": "top | bottom | shoes | outerwear | accessory | dress",
  "type": ["string"],
  "color": ["string"],
  "pattern": "describe the pattern or print 'none' if no pattern",
  "style": "describe the style or print 'unknown' if unsure",
  "fit": "[oversized | slim | regular | unknown]",
  "season": ["spring", "summer", "fall", "winter", "all"],
  "formality": "from 0.0 to 1.0 choose the formality level (0 = very casual, 1 = very formal)",
  "description": "describe the clothing item, including details about design, texture, and any other notable features."
}

Rules:
- no markdown
- JSON only
- if unsure -> "unknown"
"""


def extract_attributes(image):
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": image,
                },
                {
                    "type": "text",
                    "text": EXTRACTION_PROMPT,
                },
            ],
        }
    ]

    inputs = get_processor().apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
        downsample_mode="16x",
        max_slice_nums=16,
    )

    inputs = {
        k: v.to(get_model().device)
        for k, v in inputs.items()
    }

    with torch.no_grad():
        generated_ids = get_model().generate(
            **inputs,
            max_new_tokens=512,
        )

    generated_ids_trimmed = [
        out_ids[len(in_ids):]
        for in_ids, out_ids in zip(
            inputs["input_ids"],
            generated_ids,
        )
    ]

    result = get_processor().batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )[0]

    print("MODEL RESPONSE:")
    print(result)

    result = result.replace("```json", "").replace("```", "").strip()

    try:
        # Use raw_decode to handle extra text after valid JSON
        start_idx = result.find('{')
        if start_idx == -1:
            return {"error": "No JSON found", "raw_response": result}
        decoder = json.JSONDecoder()
        parsed, _ = decoder.raw_decode(result, start_idx)
        return parsed
    except Exception:
        return {
            "error": "Invalid JSON",
            "raw_response": result
        }