import os
import io
import threading
import torch

from PIL import Image
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import JSONResponse

from unsloth import FastVisionModel
from transformers import TextIteratorStreamer
from groq import Groq

MODEL_REPO = "Rady10/Plant-Disease-Qwen3VL-2B"

MAX_SEQ_LENGTH = 512

VISION_MAX_NEW_TOKENS = 256

SUMMARY_MAX_TOKENS = 220

GROQ_API_KEY = os.environ["GROQ_API_KEY"]

GROQ_MODEL = "llama-3.3-70b-versatile"

print("Loading vision model...")

model, tokenizer = FastVisionModel.from_pretrained(
    MODEL_REPO,
    load_in_4bit=True,
    max_seq_length=MAX_SEQ_LENGTH,
    use_gradient_checkpointing="unsloth",
)

FastVisionModel.for_inference(model)

print("Vision model loaded.")

groq_client = Groq(api_key=GROQ_API_KEY)

VISION_SYSTEM_PROMPT = """
You are a professional plant disease expert.

You identify diseases from plant images and answer agriculture questions.
"""

SUMMARY_SYSTEM_PROMPT = """
You are an agriculture assistant.

Rewrite the diagnosis in simple Egyptian Arabic.

Rules:
- 3-6 short sentences.
- Mention disease.
- Mention likely cause.
- Mention treatment.
- Do not invent information.
"""

def run_vision_model(image, question):

    messages = [
        {
            "role": "system",
            "content": [
                {
                    "type":"text",
                    "text":VISION_SYSTEM_PROMPT
                }
            ]
        },
        {
            "role":"user",
            "content":[
                {
                    "type":"image",
                    "image":image
                },
                {
                    "type":"text",
                    "text":question or "Diagnose this plant."
                }
            ]
        }
    ]

    prompt = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True
    )

    inputs = tokenizer(
        image,
        prompt,
        add_special_tokens=False,
        return_tensors="pt"
    )

    if torch.cuda.is_available():
        inputs = inputs.to("cuda")

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=VISION_MAX_NEW_TOKENS,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.15,
    )

    thread = threading.Thread(
        target=model.generate,
        kwargs=kwargs,
    )

    thread.start()

    text = ""

    for piece in streamer:
        text += piece

    return text.strip()

def summarize_with_groq(raw, question):

    response = groq_client.chat.completions.create(
        model=GROQ_MODEL,
        messages=[
            {
                "role":"system",
                "content":SUMMARY_SYSTEM_PROMPT
            },
            {
                "role":"user",
                "content":f"""
Question:
{question}

Diagnosis:

{raw}
"""
            }
        ],
        max_tokens=SUMMARY_MAX_TOKENS,
        temperature=0.6,
    )

    return response.choices[0].message.content

app = FastAPI(title="Plant Disease API")

@app.post("/analyze")
async def analyze(
    image: UploadFile = File(...),
    question: str = Form("")
):

    image_bytes = await image.read()

    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

    raw = run_vision_model(image, question)

    summary = summarize_with_groq(raw, question)

    return JSONResponse(
        {
            "question": question,
            "summary": summary,
            "technical_diagnosis": raw
        }
    )