import os
import sys
from flask import Flask, request, jsonify
from flask_cors import CORS # Keep this import, but we'll use it differently for now
import numpy as np 
import json 
import google.api_core.exceptions

from dotenv import load_dotenv

import google.generativeai as genai
from google.generativeai.types import GenerationConfig

print("--- Script Start: app.py ---")

# Load environment variables for local testing (Hugging Face handles secrets directly)
load_dotenv()

app = Flask(__name__)

# --- AGGRESSIVE CORS CONFIGURATION ---
# This attempts to ensure CORS headers are *always* sent.
# Define your allowed origin explicitly.
ALLOWED_ORIGIN = "https://sales-doc.vercel.app"

@app.before_request
def handle_options_requests():
    """Handle CORS preflight OPTIONS requests."""
    if request.method == 'OPTIONS':
        response = app.make_response('')
        response.headers.add('Access-Control-Allow-Origin', ALLOWED_ORIGIN)
        response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
        response.headers.add('Access-Control-Allow-Methods', 'GET,POST,PUT,DELETE,OPTIONS')
        response.headers.add('Access-Control-Allow-Credentials', 'true') # If your frontend sends credentials
        response.headers.add('Access-Control-Max-Age', '86400') # Cache preflight for 24 hours
        print(f"DEBUG: Handling OPTIONS preflight request from {request.origin}")
        print(f"DEBUG: Setting CORS headers for OPTIONS: {response.headers}")
        return response

@app.after_request
def add_cors_headers(response):
    """Add CORS headers to all responses."""
    response.headers.add('Access-Control-Allow-Origin', ALLOWED_ORIGIN)
    response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
    response.headers.add('Access-Control-Allow-Methods', 'GET,POST,PUT,DELETE,OPTIONS')
    response.headers.add('Access-Control-Allow-Credentials', 'true')
    print(f"DEBUG: Adding CORS headers to response for origin {request.origin if request.origin else 'N/A'}")
    print(f"DEBUG: Response headers: {response.headers}")
    return response

# You can comment out the Flask-CORS extension initialization if using manual headers,
# or keep it if it's not causing conflicts. For this aggressive approach, we'll
# rely on the manual headers.
# CORS(app, resources={r"/*": {"origins": ALLOWED_ORIGIN, "allow_headers": ["Content-Type", "Authorization"]}})


# --- Global Model Instances ---
sales_agent = None
gemini_model = None 
gemini_api_key_status = "Not Set" # Track API key status for logs

# --- Configure API Keys & Initialize Models ---
print("\n--- Starting API Key and Model Initialization ---")

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

if GEMINI_API_KEY:
    try:
        genai.configure(api_key=GEMINI_API_KEY)
        gemini_api_key_status = "Configured"
        print("Gemini API Key detected and configured.")
    except Exception as e:
        gemini_api_key_status = f"Configuration Failed: {e}"
        print(f"ERROR: Failed to configure Gemini API: {e}")
else:
    print("WARNING: GEMINI_API_KEY environment variable not found. Gemini LLM features will be disabled.")
    gemini_api_key_status = "Missing"

# --- DEEPMOST IMPORT FIX ---
try:
    from deepmost import sales
    print("Debug Point: Successfully imported deepmost.sales module.")
except ImportError as e:
    print(f"CRITICAL ERROR: Failed to import deepmost.sales module: {e}")
    print("This means the 'deepmost' library is not correctly installed or its path is wrong.")
    print("SalesRLAgent core model functionality will be disabled.")
    sales = None # Set sales to None if import fails, to prevent NameError later

# DeepMost SalesRLAgent Core Model Initialization
print("Debug Point: Attempting to instantiate sales.Agent (core RL model).")
if sales is not None: 
    try:
        # Relying on Dockerfile to make /.deepmost writable 
        sales_agent = sales.Agent(
            model_path="https://huggingface.co/DeepMostInnovations/sales-conversion-model-reinf-learning/resolve/main/sales_conversion_model.zip",
            auto_download=True,
            use_gpu=False      
        )
        if sales_agent is not None:
            print("Debug Point: DeepMost SalesRLAgent core model initialized successfully.")
        else:
            print("ERROR: DeepMost SalesRLAgent core model failed to initialize after constructor call (returned None).")
    except Exception as e:
        print(f"CRITICAL ERROR: DeepMost SalesRLAgent core model loading or instantiation failed.")
        print(f"Error Type: {type(e).__name__}")
        print(f"Error Message: {e}")
        import traceback
        traceback.print_exc()
        sales_agent = None
        print("DeepMost model initialization set to None due to error.")
else:
    print("DeepMost SalesRLAgent core model instantiation skipped because 'sales' module could not be imported.")


# Gemini LLM (1.5 Flash) Initialization
print("\nDebug Point: Attempting to initialize Gemini 1.5 Flash model.")
if GEMINI_API_KEY:
    try:
        gemini_model = genai.GenerativeModel('gemini-1.5-flash-latest') 
        test_response = gemini_model.generate_content("Hello.", generation_config=GenerationConfig(max_output_tokens=10))
        print(f"Debug Point: Gemini 1.5 Flash test response: {test_response.text[:50]}...")
        print("Debug Point: Gemini LLM (1.5 Flash) initialized successfully.")
    except Exception as e:
        print(f"CRITICAL ERROR: Gemini LLM (1.5 Flash) initialization failed.")
        print(f"Error Type: {type(e).__name__}") 
        print(f"Error Message: {e}")
        print("Ensure your GEMINI_API_KEY is correct and has access to Gemini 1.5 Flash.")
        import traceback
        traceback.print_exc()
        gemini_model = None 
    print(f"Gemini Model Status: {'Initialized' if gemini_model else 'Failed to Initialize'}")
else:
    print("Debug Point: Skipping Gemini LLM initialization because GEMINI_API_KEY is not set.")
    print("Gemini Model Status: Disabled (API Key Missing)")

print("--- Finished Model Initialization Block ---\n")


# --- Flask Routes (API Endpoints only) ---
@app.route('/analyze_conversation', methods=['POST'])
def analyze_conversation():
    if sales_agent is None:
        print("ERROR: API call received for analyze_conversation but sales_agent (core) is None.")
        return jsonify({"error": "SalesRLAgent core model not initialized on backend. Check Space logs for DeepMost initialization errors."}), 500

    try:
        data = request.get_json()
        if not data or 'conversation' not in data:
            return jsonify({"error": "Invalid request. 'conversation' field is required."}), 400

        conversation = data['conversation']
        if not isinstance(conversation, list) or not all(isinstance(turn, str) for turn in conversation):
            return jsonify({"error": "'conversation' must be a list of strings."}), 400

        print(f"Processing /analyze_conversation for: {conversation}")

        all_analysis_results = []
        full_conversation_so_far = []

        for i, turn_message in enumerate(conversation):
            full_conversation_so_far.append(turn_message)
            
            deepmost_analysis = sales_agent.analyze_conversation_progression(full_conversation_so_far, print_results=False)
            
            probability = 0.0
            if deepmost_analysis and len(deepmost_analysis) > 0:
                probability = deepmost_analysis[-1]['probability']
            
            llm_metrics = {} 
            llm_per_turn_suggestion = ""

            turn_result = {
                "turn": i + 1,
                "speaker": turn_message.split(":")[0].strip() if ":" in turn_message else "Unknown",
                "message": turn_message,
                "probability": probability,
                "status": "calculated", 
                "metrics": llm_metrics, 
                "llm_per_turn_suggestion": llm_per_turn_suggestion 
            }
            all_analysis_results.append(turn_result)
        
        print(f"Successfully processed /analyze_conversation. Returning {len(all_analysis_results)} results.")
        return jsonify({"results": all_analysis_results, "llm_advice_pending": True}), 200

    except Exception as e:
        print(f"ERROR: Exception during /analyze_conversation: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": f"An error occurred during analysis: {str(e)}"}), 500

@app.route('/get_llm_advice', methods=['POST'])
def get_llm_advice():
    if gemini_model is None:
        print("ERROR: LLM advice requested but Gemini LLM is not initialized or available.")
        return jsonify({"points": ["LLM advice unavailable. Gemini failed to load on backend. Check Space logs or add GEMINI_API_KEY secret."]}), 500

    try:
        data = request.get_json()
        conversation = data.get('conversation', [])
        if not conversation:
            return jsonify({"points": ["No conversation provided for LLM advice."]}), 400

        full_convo_text = "\n".join(conversation)
        advice_prompt = (
            f"Analyze the entire following sales conversation:\n\n"
            f"{full_convo_text}\n\n"
            f"As a concise sales coach, provide actionable advice to the salesperson on how to best progress this sales call towards a successful outcome. "
            f"Provide this advice as a JSON object with a single key 'points' which is an array of strings, where each string is a distinct, actionable bullet point. "
            f"Do NOT include any other text outside the JSON object. Ensure the JSON is well-formed and complete."
        )
        print(f"Processing /get_llm_advice. Prompting Gemini: {advice_prompt[:200]}...") 
        try:
            gemini_response = gemini_model.generate_content(
                [advice_prompt],
                generation_config=GenerationConfig(
                    response_mime_type="application/json",
                    response_schema={"type": "OBJECT", "properties": {"points": {"type": "ARRAY", "items": {"type": "STRING"}}}, "required": ["points"]},
                    max_output_tokens=300,
                    temperature=0.6
                )
            )
            raw_json_string = ""
            if gemini_response and gemini_response.candidates and len(gemini_response.candidates) > 0 and \
               gemini_response.candidates[0].content and gemini_response.candidates[0].content.parts and \
               len(gemini_response.candidates[0].content.parts) > 0:
                raw_json_string = gemini_response.candidates[0].content.parts[0].text.strip()
                print(f"Raw LLM JSON response: {raw_json_string}")
            else:
                print("WARNING: Empty or malformed LLM response for overall advice.")
                return jsonify({"points": ["LLM returned an empty or malformed response. Try again or check conversation length."]}), 200

            parsed_advice = {}
            try:
                parsed_advice = json.loads(raw_json_string)
                if "points" in parsed_advice and isinstance(parsed_advice["points"], list):
                    print(f"Successfully parsed Gemini advice: {parsed_advice}")
                    return jsonify(parsed_advice), 200
                else:
                    print(f"WARNING: LLM did not return 'points' array in structured advice: {raw_json_string}")
                    return jsonify({"points": ["LLM response was not structured as expected (missing 'points' array). Raw: " + raw_json_string[:100] + "..."]}), 200
            except json.JSONDecodeError as json_e:
                print(f"ERROR: JSON parsing error for overall advice: {json_e}. Raw string: {raw_json_string}")
                return jsonify({"points": ["Error parsing LLM JSON advice. This happens with incomplete LLM responses (e.g., due to API rate limits or max tokens). Please try a shorter conversation or wait a moment. Raw response starts with: " + raw_json_string[:100] + "..."]})
            except Exception as parse_e:
                print(f"ERROR: General error during JSON parsing attempt for chat_llm (Gemini): {parse_e}. Raw string: {raw_json_string}")
                return jsonify({"points": ["General error with LLM JSON parsing. Raw response starts with: " + raw_json_string[:100] + "..."]})
                
        except google.api_core.exceptions.ResourceExhausted as quota_e:
            print(f"ERROR: Quota Exceeded for LLM advice: {quota_e}")
            return jsonify({"points": ["Quota Exceeded: Cannot generate overall LLM advice due to API rate limits. Please try again in a minute or two."]}), 200
        except Exception as e:
            print(f"ERROR: Exception generating structured Gemini advice: {e}")
            import traceback
            traceback.print_exc()
            return jsonify({"points": [f"Error generating LLM advice: {type(e).__name__} - {e}"]}), 200

    except Exception as e:
        print(f"ERROR: An unexpected error occurred in /get_llm_advice: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({"points": [f"An unexpected error occurred: {type(e).__name__} - {e}"]}), 500


@app.route('/chat_llm', methods=['POST'])
def chat_llm():
    if gemini_model is None:
        print("ERROR: Gemini LLM instance is not initialized or available for chat.")
        return jsonify({"error": "LLM chat functionality unavailable. Gemini failed to load."}), 500
    
    try:
        data = request.get_json()
        user_message = data.get('message', '')
        if not user_message:
            return jsonify({"error": "No message provided."}), 400

        print(f"Processing /chat_llm. Received message: {user_message}")

        general_chat_prompt = f"Respond to the following message concisely: '{user_message}'"
        chat_response_obj = gemini_model.generate_content(
            general_chat_prompt,
            generation_config=GenerationConfig(max_output_tokens=150, temperature=0.7)
        )
        chat_response = chat_response_obj.text.strip()
        print(f"Gemini Raw Chat Response: {chat_response}")

        json_prompt = (
            f"Analyze the following message: '{user_message}'. "
            f"Provide a JSON object with 'summary', 'sentiment' (positive/neutral/negative), "
            f"and 'keywords' (array of strings). Do not include any other text outside the JSON block."
        )
        json_response_obj = gemini_model.generate_content(
            [json_prompt],
            generation_config=GenerationConfig(
                response_mime_type="application/json",
                max_output_tokens=200,
                temperature=0.1
            )
        )
        json_response = json_response_obj.text.strip()
        print(f"Gemini Raw JSON Prompt Response: {json_response}")
        
        parsed_json_output = None
        try:
            parsed_json_output = json.loads(json_response)
            print(f"Parsed JSON from Gemini chat: {parsed_json_output}")

        except json.JSONDecodeError as e:
            print(f"ERROR: JSON parsing error for chat_llm (Gemini): {e}. Raw string: {json_response}")
        except Exception as e:
            print(f"ERROR: General error during JSON parsing attempt for chat_llm (Gemini): {e}. Raw string: {json_response}")

        return jsonify({
            "user_message": user_message,
            "raw_chat_response": chat_response,
            "raw_json_prompt_response": json_response,
            "parsed_json_metrics": parsed_json_output,
            "status": "success"
        }), 200

    except Exception as e:
        print(f"ERROR: Error during LLM chat: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({"error": f"An error occurred during LLM chat: {str(e)}"}), 500

# Health check endpoint for Hugging Face Spaces (optional, but good practice)
@app.route('/health', methods=['GET'])
def health_check():
    status = {
        "status": "up",
        "deepmost_model_initialized": sales_agent is not None,
        "gemini_llm_initialized": gemini_model is not None,
        "gemini_api_key_status": gemini_api_key_status,
        "message": "Application is running"
    }
    # Provide more detail if a component failed
    if sales_agent is None:
        status["message"] = "Application running, but DeepMost model failed to initialize."
        status["status"] = "degraded"
    if gemini_model is None and gemini_api_key_status != "Missing": # Only degraded if API key was provided but init failed
         status["message"] = "Application running, but Gemini LLM failed to initialize."
         status["status"] = "degraded"
    elif gemini_model is None and gemini_api_key_status == "Missing":
        status["message"] = "Application running. Gemini LLM disabled (no API key)."


    print(f"Health check requested. Status: {status}")
    return jsonify(status), 200

# --- Main Execution Block ---
if __name__ == '__main__':
    try:
        print("Attempting to start Flask app (this block is primarily for local execution).")
        print("Application setup complete. Expecting Gunicorn to take over.")
        
    except Exception as startup_exception:
        print(f"CRITICAL: An unhandled exception occurred during Flask app setup: {startup_exception}")
        import traceback
        traceback.print_exc()
        sys.exit(1) # Exit with error code if startup fails