from flask import Flask, request, Response, jsonify, stream_with_context
from flask_cors import CORS
import os, json, httpx
import config
import base64
from io import BytesIO
import io
from PIL import Image
import requests


app = Flask(__name__)
CORS(app)

OPENROUTER_API_KEY = os.getenv("API_KEYS")

@app.get("/health")
def health():
    return jsonify({"status": "ready"}), 200

def _headers():
    return {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://YOURSPACE.hf.space",
        "X-Title": "Minimal Streaming Test",
    }

def upload_image_to_imgbb(image_bytes, filename):
    """Upload image to imgbb with 5 minute expiration and return public URL"""
    IMGBB_API_KEY = "22f355dd0069ef8742dc94f0c80e4c4d"
    
    try:
        # Convert bytes to base64 for imgbb upload
        base64_image = base64.b64encode(image_bytes).decode('utf-8')
        
        response = requests.post(
            "https://api.imgbb.com/1/upload",
            data={
                "key": IMGBB_API_KEY,
                "image": base64_image,
                "name": filename,
                "expiration": 300  # 300 seconds = 5 minutes
            },
            timeout=10
        )
        
        if response.status_code == 200:
            return response.json()['data']['url']
        else:
            print(f"ImgBB error: {response.text}")
            return None
    except Exception as e:
        print(f"Upload error: {e}")
        return None

# Fixed _payload: For images, use data[1] directly as the public URL (no prefixing with app URL, since it's now hosted on imgbb).
def _payload(q: str):
    data=q.split("%")
    content=""
    system=""
    if data[0]=="Text":
        content=data[1]
    else:
        # For images, content must be an array with text + image
        content=[
          {
            "type": "text",
            "text": "Analyze image and get question"
          },
          {
            "type": "image_url",
            "image_url": {
              "url": data[1]  # Direct public URL from imgbb
            }
          }
        ]
    info=config.config
    system=info[data[2]][data[3]]
    
    
    return {
        "model": "meta-llama/llama-4-maverick-17b-128e-instruct",
        "stream": True,
        "messages": [{"role":"system","content":system},{"role": "user", "content": content}],
        "temperature": 0.7,
        "top_p":0.9,
        "max_tokens": 999,
    }

def stream_openrouter(q: str):
    if not OPENROUTER_API_KEY:
        yield "data: Error: Missing API key\n\n"
        return
        
    try:
        with httpx.Client(timeout=30.0) as client:
            with client.stream("POST", 
                             "https://api.groq.com/openai/v1/chat/completions",
                             headers=_headers(),
                             json=_payload(q)) as response:
                
                response.raise_for_status()
                
                for line in response.iter_lines():
                    if line.startswith('data: '):
                        data = line[6:]  # Remove 'data: ' prefix
                        if data == '[DONE]':
                            break
                        try:
                            json_data = json.loads(data)
                            if 'choices' in json_data:
                                content = json_data['choices'][0].get('delta', {}).get('content')
                                print(content)
                                if content:
                                    yield f"data: {json.dumps(content)}\n\n"
                        except json.JSONDecodeError:
                            continue
                
                yield "data: [DONE]\n\n"
                
    except httpx.HTTPError as e:
        yield f"data: Error: {str(e)}\n\n"
    except Exception as e:
        yield f"data: Unexpected error: {str(e)}\n\n"

@app.post("/ask")
def ask():
    payload = request.get_json(silent=True) or {}
    q = payload.get("question", "").strip()
    lang = payload.get("selectedLanguage", "").strip()
    mode = payload.get("selectedMode", "").strip()
    prompt=""
    q2lang=""
    if not q:
        q = "Say hello in one short sentence."

    prompt="Text"+"%"+q+"%"+mode+"%"+lang
    print(prompt)
        
    return Response(
        stream_with_context(stream_openrouter(prompt)),
        mimetype="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no"
        }
    )

# Fixed /askimage endpoint: 
# - Use upload_image_to_imgbb to get a public URL instead of saving to local disk (which isn't served statically).
# - Removed invalid 'if not image_url:' check (path is always truthy).
# - Removed local file save.
# - Updated prompt construction to use the public URL directly.
@app.post("/askimage")
def askimage():
    # Get form data
    image_file = request.files.get('image')
    lang = request.form.get('selectedLanguage', '').strip()
    mode = request.form.get('selectedMode', '').strip()
    
    if not image_file:
        return jsonify({"error": "No image provided"}), 400
    
    try:
        # Read image bytes
        image_bytes = image_file.read()
        
        
        # Load image from bytes
        img = Image.open(io.BytesIO(image_bytes))
        
        # Resize if larger than 1024x1024 to reduce size (adjust as needed)
        max_size = 1024
        if img.size[0] > max_size or img.size[1] > max_size:
            img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
        
        # Save compressed version (quality 85 for balance of speed/size)
        output = io.BytesIO()
        img.save(output, format='JPEG', quality=85, optimize=True)
        compressed_bytes = output.getvalue()
        
        # Use compressed bytes for upload
        image_url = upload_image_to_imgbb(compressed_bytes, image_file.filename)
        
        if not image_url:
            return jsonify({"error": "Failed to upload image to host"}), 500
        
        # Build prompt with public URL
        prompt = "Image"+"%"+image_url+"%"+mode+"%"+lang
        print(f"Image URL: {image_url}")
        
        return Response(
            stream_with_context(stream_openrouter(prompt)),
            mimetype="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "X-Accel-Buffering": "no"
            }
        )
    except Exception as e:
        return jsonify({"error": str(e)}), 500

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