import json
import base64
import requests
from pathlib import Path
from typing import Optional, List, Dict, Any
from flask import Flask, request, jsonify

app = Flask(__name__)

ROBOFLOW_URL = "https://playground.roboflow.com/api/infer/openrouter/mistralai%2Fmistral-medium-3.1"
DEFAULT_MODEL = "mistralai/mistral-medium-3.1"

def encode_image_to_base64_from_path(image_path: Optional[str]) -> str:
    if not image_path:
        return ""
    p = Path(image_path)
    if not p.is_file():
        raise FileNotFoundError(f"Image not found: {image_path}")
    return base64.b64encode(p.read_bytes()).decode("utf-8")

def encode_image_to_base64_from_url(image_url: str, timeout: int = 30) -> str:
    resp = requests.get(image_url, timeout=timeout)
    resp.raise_for_status()
    return base64.b64encode(resp.content).decode("utf-8")

def normalize_messages_to_prompt(messages: List[Dict[str, str]]) -> str:
    # Convert structured chat history into a single prompt string
    # Expected roles: system, user, assistant
    # Output example:
    # System: ...
    # User: ...
    # Assistant: ...
    lines = []
    for msg in messages:
        role = str(msg.get("role", "")).strip().lower()
        content = str(msg.get("content", "")).strip()
        if not content:
            continue
        if role == "system":
            prefix = "System"
        elif role == "assistant":
            prefix = "Assistant"
        else:
            prefix = "User"
        lines.append(f"{prefix}: {content}")
    # Optionally add a final cue for the assistant to respond
    return "\n".join(lines).strip()

def extract_output(obj: dict) -> str:
    # Try common shapes:
    # 1) { data: { outputs: [ { output: "..." } ] } }
    try:
        out = obj["data"]["outputs"][0]["output"]
        if isinstance(out, str):
            return out
    except Exception:
        pass
    # 2) Flat { "output": "..." }
    if isinstance(obj, dict) and isinstance(obj.get("output"), str):
        return obj["output"]
    # Fallback: empty
    return ""

def call_roboflow_playground(image_b64: str, prompt: str, model: str = DEFAULT_MODEL, timeout: int = 60) -> Dict[str, Any]:
    payload = {
        "inputs": {
            "image": {"type": "base64", "value": image_b64},
            "prompt": prompt,
        },
        "model": model,
        "stream": False,
    }
    headers = {
        "Content-Type": "application/json",
        "Origin": "https://playground.roboflow.com",
        "Referer": "https://playground.roboflow.com/open-prompt",
        "User-Agent": "python-requests/2.x",
        "Accept": "application/json",
    }
    resp = requests.post(ROBOFLOW_URL, headers=headers, data=json.dumps(payload), timeout=timeout)
    resp.raise_for_status()
    return resp.json()

@app.route("/infer", methods=["POST"])
def infer():
    try:
        data = request.get_json(force=True, silent=False) or {}

        # Image sources: image_base64, image_url, or image_path (optional local)
        image_b64 = (data.get("image_base64") or "").strip()
        image_url = (data.get("image_url") or "").strip()
        image_path = (data.get("image_path") or "").strip()

        if not image_b64:
            if image_url:
                try:
                    image_b64 = encode_image_to_base64_from_url(image_url)
                except Exception as e:
                    return jsonify({"error": f"Failed to fetch/encode image_url: {str(e)}"}), 400
            elif image_path:
                try:
                    image_b64 = encode_image_to_base64_from_path(image_path)
                except Exception as e:
                    return jsonify({"error": f"Failed to read/encode image_path: {str(e)}"}), 400
            else:
                return jsonify({"error": "Provide one of: image_base64, image_url, or image_path"}), 400

        # Messages to prompt
        messages = data.get("messages", [])
        if not isinstance(messages, list) or not all(isinstance(m, dict) for m in messages):
            return jsonify({"error": "messages must be a list of {role, content}"}), 400

        prompt = normalize_messages_to_prompt(messages)
        if not prompt:
            return jsonify({"error": "messages is empty or has no valid content"}), 400

        # Model (optional override)
        model = (data.get("model") or DEFAULT_MODEL).strip() or DEFAULT_MODEL

        # Call external API
        raw_response = call_roboflow_playground(image_b64=image_b64, prompt=prompt, model=model)

        # Extract simplified output
        output_text = extract_output(raw_response)

        return jsonify({
            "output": output_text,
            "raw_response": raw_response
        }), 200

    except requests.HTTPError as e:
        status = getattr(e.response, "status_code", 502)
        text = getattr(e.response, "text", "")
        return jsonify({"error": "HTTP error from upstream", "status": status, "response_text": text}), 502
    except Exception as e:
        return jsonify({"error": str(e)}), 500

def main():
    # Run Flask on port 7860
    app.run(host="0.0.0.0", port=7860, debug=False)

if __name__ == "__main__":
    main()