import os

import spaces
import gradio as gr

MODEL_ID = "zieglerd/RussianConcsExt"
MAX_INPUT_TOKENS = 16_384
IS_ZERO_GPU = os.getenv("SPACES_ZERO_GPU") == "1"

tokenizer = None
model = None
load_error = None

if IS_ZERO_GPU:
    try:
        import torch
        from peft import AutoPeftModelForCausalLM
        from transformers import AutoTokenizer, BitsAndBytesConfig

        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
        )
        model = AutoPeftModelForCausalLM.from_pretrained(
            MODEL_ID,
            quantization_config=quantization_config,
            device_map="cuda",
            dtype=torch.bfloat16,
            attn_implementation="sdpa",
        ).eval()
    except Exception as exc:
        load_error = f"{type(exc).__name__}: {exc}"
        print(f"Model loading failed: {load_error}", flush=True)


EXAMPLE_CONTRACT = """ДОГОВОР ПОСТАВКИ №458/26

ООО «Поставщик» обязуется поставить ООО «Заказчик»
общехозяйственные товары согласно спецификации.

Стоимость договора составляет 34 132 686,08 рублей,
в том числе НДС 5%.

Оплата производится следующим образом:
85% — авансовый платеж.
15% — после поставки товара."""


def _extract_json(text: str):
    """Parse the first complete JSON object from model output when possible."""
    import json

    cleaned = text.strip()
    if cleaned.startswith("```"):
        lines = cleaned.splitlines()
        if lines and lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        cleaned = "\n".join(lines).strip()

    try:
        return json.loads(cleaned)
    except json.JSONDecodeError:
        pass

    start = cleaned.find("{")
    if start < 0:
        return None

    depth = 0
    in_string = False
    escaped = False
    for index, char in enumerate(cleaned[start:], start=start):
        if escaped:
            escaped = False
            continue
        if char == "\\" and in_string:
            escaped = True
            continue
        if char == '"':
            in_string = not in_string
        elif not in_string:
            if char == "{":
                depth += 1
            elif char == "}":
                depth -= 1
                if depth == 0:
                    try:
                        return json.loads(cleaned[start : index + 1])
                    except json.JSONDecodeError:
                        return None
    return None


@spaces.GPU(duration=180, size="large")
def extract_contract(contract_text: str, max_new_tokens: int):
    """Extract structured contract data from Russian text and return JSON."""
    if not contract_text or not contract_text.strip():
        return {}, "", "Введите текст договора."
    if load_error:
        return {}, "", f"Не удалось загрузить модель: {load_error}"
    if model is None or tokenizer is None:
        return (
            {},
            "",
            "Демо ожидает предоставления ZeroGPU. Код и интерфейс уже развернуты.",
        )

    messages = [{"role": "user", "content": contract_text.strip()}]
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=MAX_INPUT_TOKENS,
    ).to("cuda")

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=False,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated_ids = output_ids[0, inputs["input_ids"].shape[1] :]
    raw_output = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
    parsed = _extract_json(raw_output)
    if parsed is None:
        return (
            {"raw_output": raw_output},
            raw_output,
            "Модель вернула ответ, но он не является корректным JSON.",
        )
    return parsed, raw_output, "Готово. Проверьте результат перед использованием."


with gr.Blocks(title="Russian Contract Extractor") as demo:
    gr.Markdown(
        """
# Извлечение данных из российских договоров

Демо LoRA-адаптера
[zieglerd/RussianConcsExt](https://huggingface.co/zieglerd/RussianConcsExt)
для Qwen2.5-32B-Instruct. Вставьте текст договора, чтобы получить
структурированный JSON.

> Результат предназначен для информационного извлечения, не является
> юридической консультацией и требует проверки.
"""
    )
    with gr.Row():
        with gr.Column(scale=3):
            contract_input = gr.Textbox(
                label="Текст договора",
                value=EXAMPLE_CONTRACT,
                lines=18,
                placeholder="Вставьте распознанный или исходный текст договора…",
            )
            max_tokens = gr.Slider(
                minimum=256,
                maximum=2048,
                value=1024,
                step=128,
                label="Максимум токенов ответа",
            )
            extract_button = gr.Button("Извлечь данные", variant="primary")
        with gr.Column(scale=2):
            json_output = gr.JSON(label="Структурированный результат")
            status_output = gr.Markdown()

    with gr.Accordion("Исходный ответ модели", open=False):
        raw_output = gr.Code(label="Ответ", language="json")

    extract_button.click(
        fn=extract_contract,
        inputs=[contract_input, max_tokens],
        outputs=[json_output, raw_output, status_output],
        api_name="extract_contract",
    )

    gr.Examples(
        examples=[[EXAMPLE_CONTRACT, 1024]],
        inputs=[contract_input, max_tokens],
        cache_examples=True,
        cache_mode="lazy",
    )


if __name__ == "__main__":
    demo.launch(mcp_server=True)