from flask import Flask, request, jsonify
import torch
from transformers import pipeline
import os
import json # <-- Humne yeh nayi library import ki hai

# Model ka naam
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

print("Attempting to load model...")
try:
    # AI Pipeline ko load karna
    pipe = pipeline(
        "text-generation",
        model=model_name,
        torch_dtype=torch.float32, # CPU ke liye float32 behtar hai
        device=-1 # CPU par run karne ke liye force karna
    )
    print("✅ Model loaded successfully!")
except Exception as e:
    print(f"❌ Error loading model: {e}")
    pipe = None

app = Flask(__name__)

# Meta-prompt ka template
meta_prompt_template = """
You are a creative food writer and an expert chef. Based on the following cooking step, you must generate a JSON object with two keys: "tip" and "description".

1.  For the "tip" key, provide a single, short, and engaging Pro-Tip or a Fun Fact related to an ingredient or technique in the step.
2.  For the "description" key, write a creative, deep, and evocative description of the cooking action, as if you are narrating a cinematic cooking show. Use sensory language (smell, sound, sight).

You must not output any text other than the final, clean JSON object.

Current Cooking Step: "{step_text}"
"""

@app.route('/generate-content', methods=['POST'])
def generate_content():
    if not pipe:
        return jsonify({"error": "Model is not available or failed to load"}), 503

    data = request.get_json()
    if not data or 'step_text' not in data:
        return jsonify({"error": "step_text is missing from the request body"}), 400
    
    step_text = data.get('step_text')

    try:
        final_prompt = meta_prompt_template.format(step_text=step_text)
        
        outputs = pipe(final_prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
        
        generated_text = outputs[0]['generated_text']
        
        # Naya, behtar tareeka JSON nikalne ka
        start_index = generated_text.find('{')
        end_index = generated_text.rfind('}') + 1
        
        if start_index != -1 and end_index != -1:
            json_part = generated_text[start_index:end_index]
            # String ko JSON mein convert karna
            response_data = json.loads(json_part)
            return jsonify(response_data)
        else:
            return jsonify({"error": "Failed to find a valid JSON object in the model's response.", "raw_response": generated_text}), 500

    except json.JSONDecodeError:
        return jsonify({"error": "Model returned a malformed JSON string.", "raw_response": generated_text}), 500
    except Exception as e:
        print(f"Error during generation: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/', methods=['GET'])
def home():
    return "Creative Writer API is running!", 200

# Hugging Face Spaces ke liye
if __name__ == '__main__':
    port = int(os.environ.get("PORT", 7860))
    app.run(host='0.0.0.0', port=port)