from flask import Flask, request, jsonify
from models import (
    load_llm, text_splitter, embedding_model, prompt_template,
    extract_json_block, save_model_state, load_model_state
)
from langchain.vectorstores import Chroma

app = Flask(__name__)

# Load the model once when the server starts
print("Loading model...")
llm = load_llm()
print("Model loaded successfully!")

@app.route('/health', methods=['GET'])
def health_check():
    return jsonify({"status": "healthy", "model_loaded": True})

@app.route('/generate_workflow', methods=['POST'])
def generate_workflow():
    try:
        data = request.get_json()
        user_prompt = data.get('prompt', '')
        use_rag = data.get('use_rag', False)
        raw_docs = data.get('raw_docs', None)

        if not user_prompt.strip():
            return jsonify({"error": "Prompt is required"}), 400

        context_block = ""
        if use_rag and raw_docs:
            # Process documents for RAG
            all_chunks = []
            for doc in raw_docs:
                all_chunks.extend(text_splitter.split_text(doc))

            docs_for_vector = [{"page_content": c} for c in all_chunks]
            vector = Chroma.from_documents(docs_for_vector, embedding=embedding_model)
            retriever = vector.as_retriever(search_kwargs={"k": 3})
            retrieved_docs = retriever.get_relevant_documents(user_prompt)
            context_block = "Context:\n" + "\n".join([d.page_content for d in retrieved_docs]) + "\n"

        # Generate workflow
        full_prompt = prompt_template.format(prompt=user_prompt, context_block=context_block)
        wrapped_prompt = f"[INST] {full_prompt} [/INST]"

        output = llm(wrapped_prompt, max_tokens=1024, stop=["</s>"])
        generated_text = output["choices"][0]["text"]
        generated_json = extract_json_block(generated_text)

        if generated_json:
            # Optionally save model state for optimization
            # save_model_state(llm, "model_state.bin")
            return jsonify({
                "success": True,
                "workflow": generated_json,
                "context_used": bool(context_block)
            })
        else:
            return jsonify({"error": "Failed to extract valid JSON from generated text"}), 400

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

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)