import os
import re
import json
import base64
import textwrap
import requests
from io import BytesIO
from PIL import Image
import cv2
import numpy as np
from flask import Flask, request, jsonify, send_file
from flask_cors import CORS

app = Flask(__name__)
CORS(app)

# ── Groq ──────────────────────────────────────────────────────────────────────
def get_client():
    from groq import Groq
    return Groq(api_key=os.environ.get("GROQ_API_KEY", ""))

# ── Wikipedia ─────────────────────────────────────────────────────────────────
def get_wiki(name):
    try:
        term = name.replace(" ", "_")
        r = requests.get(
            f"https://en.wikipedia.org/api/rest_v1/page/summary/{term}",
            headers={"User-Agent": "PlantLens/1.0"},
            timeout=6
        )
        d = r.json()
        summary = d.get("extract", "")
        summary = re.sub(r'\[.*?\]', '', summary)
        summary = re.sub(r'\s{2,}', ' ', summary).strip()[:500]
        url = d.get("content_urls", {}).get("desktop", {}).get("page", "")
        return summary, url
    except Exception:
        return "", ""

# ── Identify plants via Groq vision ──────────────────────────────────────────
def identify_plants(image_bytes):
    client = get_client()
    b64 = base64.b64encode(image_bytes).decode("utf-8")

    prompt = textwrap.dedent("""\
        You are an expert botanist. Look at this image and identify EVERY plant visible.
        Reply with ONLY a JSON array, no markdown, no explanation:
        [
          {
            "common_name": "...",
            "scientific_name": "...",
            "family": "...",
            "confidence": "high|medium|low",
            "key_features": ["...", "...", "..."],
            "wikipedia_search_term": "...",
            "bbox": {"x_pct": 10, "y_pct": 10, "w_pct": 80, "h_pct": 80}
          }
        ]
        bbox values are percentages (0-100) of image width/height.
        If no plant found, return [].
    """)

    resp = client.chat.completions.create(
        model="meta-llama/llama-4-scout-17b-16e-instruct",
        messages=[{
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
                {"type": "text", "text": prompt}
            ]
        }],
        temperature=0.2,
        max_tokens=1500
    )

    raw = resp.choices[0].message.content
    cleaned = re.sub(r'```json|```', '', raw).strip()
    try:
        result = json.loads(cleaned)
        return result if isinstance(result, list) else []
    except Exception:
        return []

# ── Draw annotations on image ─────────────────────────────────────────────────
BG    = (247, 243, 238)
DOT_C = (50,  50,  50)

def annotate(image_bytes, plants):
    arr  = np.frombuffer(image_bytes, np.uint8)
    orig = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    OH, OW = orig.shape[:2]

    PAD   = max(20, OW // 30)
    W     = OW + PAD * 2
    H     = OH

    bg_bgr = (BG[2], BG[1], BG[0])
    canvas = np.full((H, W, 3), bg_bgr, dtype=np.uint8)
    canvas[0:OH, PAD:PAD+OW] = orig

    DOT_R  = max(12, OW // 65)
    sc_num = max(0.32, OW / 2400)
    FONT   = cv2.FONT_HERSHEY_SIMPLEX

    dot_positions = []

    for i, p in enumerate(plants):
        bb = p.get("bbox", {})
        cx = PAD + int((bb.get("x_pct", 50) + bb.get("w_pct", 10) / 2) / 100 * OW)
        cy =       int((bb.get("y_pct", 50) + bb.get("h_pct", 10) / 2) / 100 * OH)
        cx = min(max(cx, PAD + DOT_R + 2), PAD + OW - DOT_R - 2)
        cy = min(max(cy, DOT_R + 2), OH - DOT_R - 2)

        # White halo
        cv2.circle(canvas, (cx, cy), DOT_R + 2, (255, 255, 255), -1, cv2.LINE_AA)
        # Dark dot
        cv2.circle(canvas, (cx, cy), DOT_R, DOT_C, -1, cv2.LINE_AA)
        # Number
        num = str(i + 1)
        (nw, nh), _ = cv2.getTextSize(num, FONT, sc_num, 1)
        cv2.putText(canvas, num, (cx - nw//2, cy + nh//2), FONT, sc_num, (255,255,255), 1, cv2.LINE_AA)

        # Store dot position as percentage of final canvas for tooltip
        dot_positions.append({
            "x_pct": round(cx / W * 100, 2),
            "y_pct": round(cy / H * 100, 2),
            "name":  p.get("common_name", "Unknown")
        })

    ok, buf = cv2.imencode(".png", canvas)
    return buf.tobytes(), dot_positions


# ── Routes ────────────────────────────────────────────────────────────────────
@app.route("/")
def index():
    return send_file("index.html")

@app.route("/health")
def health():
    return jsonify({"status": "ok"})

@app.route("/analyze", methods=["POST"])
def analyze():
    if "file" not in request.files:
        return jsonify({"error": "No file"}), 400

    raw_bytes = request.files["file"].read()

    try:
        plants = identify_plants(raw_bytes)

        # Enrich with Wikipedia
        for p in plants:
            term = p.get("wikipedia_search_term") or p.get("common_name", "")
            summary, url = get_wiki(term)
            p["wiki_summary"] = summary
            p["wiki_url"]     = url

        # Annotate image
        annotated_bytes, dot_positions = annotate(raw_bytes, plants)
        annotated_b64   = base64.b64encode(annotated_bytes).decode("utf-8")

        return jsonify({
            "plants":           plants,
            "count":            len(plants),
            "annotated_image":  f"data:image/png;base64,{annotated_b64}",
            "dot_positions":    dot_positions
        })

    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    print("Starting PlantLens on port 7860...")
    app.run(host="0.0.0.0", port=7860, debug=False)