from fastapi import FastAPI, HTTPException
from models import load_model
from utils import get_prompt, preprocess_description, DiagramRequest, DiagramResponse
from utils_pic import get_pic_prompt
import re

app = FastAPI(title="Diagram → PlantUML (CPU + Qwen2.5-7B GGUF)", version="0.1.0")

llm = None

@app.on_event("startup")
def startup_event():
    global llm
    llm = load_model()  # загружаем при старте — увидим ошибки сразу

@app.get("/health")
async def health_check():
    return {
        "status": "healthy" if llm is not None else "model loading failed",
        "model_loaded": llm is not None
    }

# ... (остальное без изменений)

@app.post("/api/v1/generate_plantuml", response_model=DiagramResponse)
async def generate_plantuml(req: DiagramRequest):
    global llm
    if llm is None:
        raise HTTPException(503, "Модель ещё не загружена")

    desc = preprocess_description(req.description)
    if not desc:
        raise HTTPException(400, "Description cannot be empty")

    messages = get_prompt(desc)  # теперь list[dict]

    output = llm.create_chat_completion(
        messages,
        max_tokens=1500,           # чуть больше — примеры длинные
        temperature=0.6,           # ниже — меньше креатива, строже синтаксис
        top_p=0.9,
        top_k=35,
        repeat_penalty=1.15,       # меньше повторений
        stop=["</s>", "<|im_end|>", "@enduml\n\n", "```"],  # стоп на конце
    )

    generated_text = output["choices"][0]["message"]["content"].strip()

    # Извлечение только кода (на всякий случай)
    start_marker = "@startuml"
    end_marker = "@enduml"
    if start_marker in generated_text and end_marker in generated_text:
        start_idx = generated_text.find(start_marker)
        end_idx = generated_text.rfind(end_marker) + len(end_marker)
        plantuml_code = generated_text[start_idx:end_idx]
    else:
        plantuml_code = generated_text  # fallback

    # Финальная чистка: убираем лишние пустые строки внутри
    plantuml_code = re.sub(r'\n\s*\n', '\n', plantuml_code).strip()

    return {"plantuml_code": plantuml_code}


from fastapi import UploadFile, File
from utils_pic import (
    PicToUmlRequest, PicToUmlResponse,
    load_yolo, preprocess_image,
    run_ocr, build_graph_from_detections,
    generate_description_from_graph
)


llm = None
yolo_model = None  # ← можно явно объявить здесь для ясности

@app.on_event("startup")
def startup_event():
    global llm, yolo_model
    llm = load_model()
    yolo_model = load_yolo()  # теперь присваиваем глобальной

@app.post("/api/v1/pic_to_uml", response_model=PicToUmlResponse)
async def pic_to_uml(
    image: UploadFile = File(...),
    is_bpmn: bool = True
):
    global yolo_model, llm  # ← обязательно, если используешь глобальные llm тоже

    if not image.content_type.startswith("image/"):
        raise HTTPException(400, "Ожидается изображение (png/jpg)")

    image_bytes = await image.read()

    img_cv, pil_img = preprocess_image(image_bytes)

    if yolo_model is None:
        raise HTTPException(503, "YOLO модель не загружена")

    # В эндпоинте (app.py) измени вызов YOLO:
    results = yolo_model(
        img_cv,
        conf=0.15,  # сильно ниже — 0.15–0.20 часто спасает
        iou=0.4,  # чуть ниже, чтобы не подавлять близкие боксы
        max_det=300,  # по умолчанию 300, но можно 1000 (если много элементов)
        imgsz=1024
    )

    ocr_text = run_ocr(pil_img)

    # в pic_to_uml:


    G, nodes = build_graph_from_detections(results, ocr_text)

    messages = get_pic_prompt(ocr_text, nodes, list(G.edges()), is_bpmn)

    output = llm.create_chat_completion(
        messages,
        max_tokens=1500,
        temperature=0.65,
        top_p=0.9,
        stop=["</s>", "<|im_end|>"],
    )

    result_text = output["choices"][0]["message"]["content"].strip()

    # Если UML — можно дополнительно обрезать до @startuml ... @enduml, как в первом эндпоинте
    if not is_bpmn:
        start = result_text.find("@startuml")
        end = result_text.rfind("@enduml")
        if start != -1 and end != -1:
           result_text = result_text[start:end + 8]

    return PicToUmlResponse(result=result_text)