import os
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

import json
import torch
import re
import numpy as np
import pandas as pd
from flask import Flask, request, jsonify, render_template, redirect, url_for, flash, session
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from docx import Document
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
import torch.nn as nn
from tqdm import tqdm
import openai
from openai import OpenAI
import time
from huggingface_hub import login
import logging
from functools import wraps
import requests
from safetensors.torch import save_file

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("app.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# Modify environment variable handling
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise ValueError("OPENAI_API_KEY environment variable is not set. Please set it in your .env file.")

os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
client = OpenAI(
    api_key=OPENAI_API_KEY,
)

# Model configuration
MODEL_NAME = 'dmis-lab_biobert-v1.1_20250316_195322'
MODIFIER_MODEL_NAME = 'dmis-lab_biobert-v1.1_modifier_20250406_160702'
OUTPUT_DIR = os.environ.get('CCTA_OUTPUT_DIR', 'Saved Models')

# Set device
device = torch.device('cuda' if torch.cuda.is_available() else 
                     'mps' if torch.backends.mps.is_available() else 
                     'cpu')
logger.info(f"Using device: {device}")

# Label mapping
label_mapping = {
    0: "0",
    1: "1",
    2: "2",
    3: "3",
    4: "4",
    5: "5",
    6: "N"
}
reverse_label_mapping = {v: k for k, v in label_mapping.items()}

# Modifier mapping
modifier_mapping = {
    0: "P1",  # Mild plaque burden (CAC 0-99)
    1: "P2",  # Moderate plaque burden (CAC 100-399)
    2: "P3",  # Severe plaque burden (CAC 400-999)
    3: "P4"   # Extensive plaque burden (CAC ≥1000)
}

# Load model
def load_model(model_name, num_labels):
    try:
        # Direct path to model
        model_path = os.path.join(OUTPUT_DIR, model_name.replace("/", "_"))
        
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"No model found at {model_path}")
        
        logger.info(f"Loading model from {model_path}")
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        
        # Check if safetensors version exists
        safetensors_path = os.path.join(model_path, "model.safetensors")
        pytorch_path = os.path.join(model_path, "pytorch_model.bin")
        
        # If PyTorch model exists but safetensors doesn't, convert it
        if os.path.exists(pytorch_path) and not os.path.exists(safetensors_path):
            logger.info(f"Converting PyTorch model to safetensors format")
            
            # Load the PyTorch model
            temp_model = AutoModelForSequenceClassification.from_pretrained(
                model_path, 
                num_labels=num_labels
            )
            
            # Extract state dict
            state_dict = temp_model.state_dict()
            
            # Save as safetensors
            save_file(state_dict, safetensors_path)
            logger.info(f"Model converted and saved to {safetensors_path}")
        
        # Load the model (will use safetensors if available)
        model = AutoModelForSequenceClassification.from_pretrained(
            model_path, 
            num_labels=num_labels
        )
        model.to(device)
        model.eval()
        return model, tokenizer
    except Exception as e:
        logger.error(f"Error loading model: {str(e)}")
        raise

# Load models and tokenizers
try:
    # Load classifier model for CADRADS score
    classifier_model, classifier_tokenizer = load_model(MODEL_NAME, len(label_mapping))
    logger.info("Classifier model and tokenizer loaded successfully")
    
    # Load modifier model
    modifier_model, modifier_tokenizer = load_model(MODIFIER_MODEL_NAME, len(modifier_mapping))
    logger.info("Modifier model and tokenizer loaded successfully")
except Exception as e:
    logger.error(f"Failed to load models: {str(e)}")
    raise

# Create uploads directory if it doesn't exist
UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads')
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

app = Flask(__name__)

# Get authentication credentials from environment variables
username = os.getenv('BASIC_AUTH_USERNAME')
password = os.getenv('BASIC_AUTH_PASSWORD')

# Check if credentials are set
if not username or not password:
    logger.error("BASIC_AUTH_USERNAME or BASIC_AUTH_PASSWORD environment variables are not set")
    raise ValueError("Authentication credentials are required. Set BASIC_AUTH_USERNAME and BASIC_AUTH_PASSWORD in environment variables")

# Use environment variable for secret key if available, otherwise generate a random one
app.config['SECRET_KEY'] = os.getenv('FLASK_SECRET_KEY', os.urandom(24))
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024

# Custom login required decorator
def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'user_id' not in session:
            return redirect(url_for('login', next=request.url))
        return f(*args, **kwargs)
    return decorated_function

history = []

def extract_text_from_docx(docx_path):
    try:
        doc = Document(docx_path)
        full_text = []
        
        # Extract text from paragraphs
        for para in doc.paragraphs:
            if para.text.strip():  # Only add non-empty paragraphs
                full_text.append(para.text)
        
        # Extract text from tables
        for table in doc.tables:
            for row in table.rows:
                row_text = []
                for cell in row.cells:
                    # Extract text from paragraphs inside cells
                    cell_text = ' '.join(p.text.strip() for p in cell.paragraphs if p.text.strip())
                    if cell_text:
                        row_text.append(cell_text)
                if row_text:
                    full_text.append(' | '.join(row_text))  # Use pipe to separate cell content
        
        extracted_text = '\n'.join(full_text)
        
        # Add a check to warn about potentially empty documents
        if len(extracted_text.strip()) < 10:  # Less than 10 characters is suspiciously small
            logger.warning(f"Very little text extracted ({len(extracted_text)} chars). Document may be corrupt or empty.")
        
        return extracted_text
    except Exception as e:
        logger.error(f"Error extracting text from document: {str(e)}")
        raise

def generate_json_output(report_text, prompt_template):
    prompt = prompt_template.replace("{report_text}", report_text)
    try:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,
            max_tokens=6500,
            top_p=0.9,
            frequency_penalty=0.2,
            presence_penalty=0,
        )
        response_text = response.choices[0].message.content
        
        cleaned_response_text = response_text.strip()
        if cleaned_response_text.startswith("```") and cleaned_response_text.endswith("```"):
            cleaned_response_text = cleaned_response_text[3:-3].strip()
        json_match = re.search(r'{.*}', cleaned_response_text, re.DOTALL)
        if json_match:
            cleaned_response_text = json_match.group(0)
        else:
            raise json.JSONDecodeError("No valid JSON object found", cleaned_response_text, 0)
        report_json = json.loads(cleaned_response_text)
        return report_json
    except json.JSONDecodeError as e:
        logger.error(f"JSON parsing error: {str(e)}")
        return {"error": "Unable to parse response"}
    except Exception as e:
        logger.error(f"Error generating JSON output: {str(e)}")
        return {"error": f"Error generating output: {str(e)}"}

def create_text_representation(report_json):
    """Create a consistent text representation for prediction, using same format as in training"""
    try:
        combined_text = (
            f"Technical Quality: {report_json['technical_quality']} | "
            f"Left Main: {report_json['left_main']} | "
            f"LAD: {report_json['lad']} | "
            f"LCX: {report_json['lcx']} | "
            f"RCA: {report_json['rca']} | "
            f"Stenosis Summary: {report_json['stenosis_summary']} | "
            f"Calcium Score: {report_json['calcium_score']} | "
            f"Cardiac Valves: {report_json['cardiac_valves']} | "
            f"Pericardium: {report_json['pericardium']} | "
            f"Dominance: {report_json['dominance']} | "
            f"Impression: {report_json['impression']}"
        )
        return combined_text
    except KeyError as e:
        logger.error(f"Missing field in report JSON: {str(e)}")
        raise

def predict_cadrads_score(report_json):
    try:
        # Create text representation consistent with training
        combined_text = create_text_representation(report_json)
        
        with torch.no_grad():
            # Predict CADRADS score with classifier model
            classifier_encoding = classifier_tokenizer(combined_text, truncation=True, padding='max_length', max_length=512, return_tensors='pt')
            classifier_input_ids = classifier_encoding['input_ids'].to(device)
            classifier_attention_mask = classifier_encoding['attention_mask'].to(device)
            classifier_outputs = classifier_model(classifier_input_ids, attention_mask=classifier_attention_mask)
            
            # Get CADRADS score (highest probability class)
            classifier_probabilities = torch.nn.functional.softmax(classifier_outputs.logits, dim=1)
            _, predicted_cadrads = torch.max(classifier_probabilities, dim=1)
            cadrads_class = predicted_cadrads.item()
            
            # Convert to label
            if cadrads_class in label_mapping:
                cadrads_score = label_mapping[cadrads_class]
            else:
                logger.warning(f"Predicted CADRADS class {cadrads_class} not in label mapping")
                cadrads_score = "Unknown"
            
            # If CADRADS score is 0 or N, don't use modifier
            if cadrads_score == "0" or cadrads_score == "N":
                logger.info(f"CADRADS score is {cadrads_score}, not displaying modifier")
                return cadrads_score, "None"
            
            # Only predict modifier if CADRADS score is not 0 or N
            modifier_encoding = modifier_tokenizer(combined_text, truncation=True, padding='max_length', max_length=512, return_tensors='pt')
            modifier_input_ids = modifier_encoding['input_ids'].to(device)
            modifier_attention_mask = modifier_encoding['attention_mask'].to(device)
            modifier_outputs = modifier_model(modifier_input_ids, attention_mask=modifier_attention_mask)
            
            # Get modifier (highest probability class)
            modifier_probabilities = torch.nn.functional.softmax(modifier_outputs.logits, dim=1)
            _, predicted_modifier = torch.max(modifier_probabilities, dim=1)
            modifier_class = predicted_modifier.item()
            
            # Convert to label
            if modifier_class in modifier_mapping:
                modifier = modifier_mapping[modifier_class]
            else:
                logger.warning(f"Predicted modifier class {modifier_class} not in modifier mapping")
                modifier = "None"
            
            return cadrads_score, modifier
    except Exception as e:
        logger.error(f"Error in prediction: {str(e)}")
        return "Error", "None"

# Load prompt template
try:
    prompt_template_path = "prompt_template_modifiers_new.md"
    with open(prompt_template_path, "r") as file:
        prompt_template = file.read()
    logger.info("Prompt template loaded successfully")
except Exception as e:
    logger.error(f"Error loading prompt template: {str(e)}")
    prompt_template = None

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        input_username = request.form['username']
        input_password = request.form['password']
        
        if input_username == username and input_password == password:
            session['user_id'] = input_username
            next_page = request.args.get('next')
            if next_page and next_page.startswith('/'):
                return redirect(next_page)
            return redirect(url_for('home'))
        else:
            flash('Invalid username or password', 'danger')
    
    return render_template('login.html')

@app.route('/logout')
def logout():
    session.pop('user_id', None)
    flash('You have been logged out', 'info')
    return redirect(url_for('login'))

@app.route('/')
@login_required
def home():
    return render_template('index.html', history=history)

@app.route('/predict', methods=['POST'])
@login_required
def predict():
    if not prompt_template:
        flash('Prompt template not loaded', 'error')
        return redirect(url_for('home'))
    
    if 'file' not in request.files:
        flash('No file selected', 'error')
        return redirect(url_for('home'))
        
    file = request.files['file']
    if file.filename == '':
        flash('No file selected', 'error')
        return redirect(url_for('home'))
        
    if not file.filename.endswith('.docx'):
        flash('Invalid file format. Please upload a .docx file.', 'error')
        return redirect(url_for('home'))

    try:
        session['processing'] = True
        filename = os.path.join(
            app.config['UPLOAD_FOLDER'],
            file.filename
        )
        
        os.makedirs(os.path.dirname(filename), exist_ok=True)
        file.save(filename)
        logger.info(f"File saved to {filename}")
        
        try:
            # Extract text from document
            report_text = extract_text_from_docx(filename)
            logger.info(f"Text extracted from {file.filename}")
            os.remove(filename)
            
            # Generate structured JSON from report text
            report_json = generate_json_output(report_text, prompt_template)
            if "error" in report_json:
                logger.error(f"Error in JSON generation: {report_json['error']}")
                flash('Error processing report', 'error')
                return redirect(url_for('home'))

            # Predict CADRADS score
            cadrads_score, modifier = predict_cadrads_score(report_json)
            
            # Log prediction
            logger.info(f"Prediction for {file.filename}: CADRADS {cadrads_score}, Modifier {modifier}")
            
            # Add to history
            report_id = len(history) + 1
            history.append({
                'report_id': report_id,
                'prediction': cadrads_score,
                'modifier': modifier,
                'text': report_text,
                'timestamp': time.strftime('%Y-%m-%d %H:%M:%S'),
                'filename': file.filename,
                'json_output': report_json  # Store the JSON output in the report history
            })

            flash('Report processed successfully!', 'success')
            return redirect(url_for('home'))
            
        except Exception as e:
            if os.path.exists(filename):
                os.remove(filename)
            logger.error(f"Error processing file content: {str(e)}")
            raise e
            
    except Exception as e:
        logger.error(f"Error in prediction process: {str(e)}")
        flash(f'Error processing file: {str(e)}', 'error')
        return redirect(url_for('home'))
        
    finally:
        session.pop('processing', None)

@app.route('/view_report/<int:report_id>')
@login_required
def view_report(report_id):
    report = next((item for item in history if item["report_id"] == report_id), None)
    if report:
        if 'mesa_results' not in report:
            report['mesa_results'] = None
        
        # Make the JSON data available to the template for extracting calcium score
        if 'json_output' in report:
            report['json_data'] = report['json_output']
        
    return render_template('view_report.html', report=report)

@app.route('/calculate_mesa_score', methods=['POST'])
@login_required
def calculate_mesa_score():
    try:
        # Get form data from request
        data = request.json
        if not data:
            return jsonify({"error": "No data provided"}), 400
            
        # Validate required fields
        required_fields = ['gender', 'age', 'ethnicity', 'diabetes', 'smoke', 
                          'fhhx', 'chol', 'hdl', 'sbp', 'lipid', 'htnmed']
        for field in required_fields:
            if field not in data:
                return jsonify({"error": f"Missing required field: {field}"}), 400
                
        # Prepare API request
        api_endpoint = "https://www.mesa-nhlbi.org/MESACHDRiskCoronaryAge/api/score"
        
        # Add CAC score if provided
        cac_param = ""
        if 'cac' in data and data['cac']:
            cac_param = f"/{data['cac']}"
        else:
            cac_param = "/-9"  # Default value for CAC if not provided
            
        # Construct full URL with parameters
        url = f"{api_endpoint}/{data['gender']}/{data['age']}/{data['ethnicity']}/{data['diabetes']}/{data['smoke']}/{data['fhhx']}/{data['chol']}/{data['hdl']}/{data['sbp']}/{data['lipid']}/{data['htnmed']}{cac_param}"
        
        # Call MESA API
        logger.info(f"Calling MESA API: {url}")
        response = requests.get(url)
        
        if response.status_code != 200:
            logger.error(f"MESA API error: {response.status_code} - {response.text}")
            return jsonify({"error": f"MESA API returned error: {response.status_code}"}), 500
            
        # Process and return the results
        result = response.json()
        logger.info(f"MESA API result: {result}")
        
        # Extract and format relevant data
        formatted_result = {
            'cac_riskscore': result.get('cac_riskscore', '-'),
            'cac_coronary_age': result.get('cac_coronary_age', '-'),
            'cac_age_diff': result.get('cac_age_diff', '-'),
            'nocac_riskscore': result.get('nocac_riskscore', '-'),
            'nocac_coronary_age': result.get('nocac_coronary_age', '-'),
            'nocac_age_diff': result.get('nocac_age_diff', '-')
        }
        
        # Store the form values along with the MESA results in the report data
        report_id = request.json.get('report_id')
        if report_id:
            for item in history:
                if item['report_id'] == int(report_id):
                    # Save both the results and the form input values
                    item['mesa_results'] = formatted_result
                    
                    # Create a copy of the form data without the report_id
                    form_data = {k: v for k, v in data.items() if k != 'report_id'}
                    item['mesa_form_data'] = form_data
                    
                    logger.info(f"MESA results and form data stored for report {report_id}")
                    break
        
        return jsonify(formatted_result)
        
    except Exception as e:
        logger.error(f"Error calculating MESA score: {str(e)}")
        return jsonify({"error": str(e)}), 500

@app.errorhandler(413)
def too_large(e):
    flash('File is too large. Maximum size is 16MB.', 'error')
    return redirect(url_for('home'))

@app.errorhandler(404)
def not_found(e):
    return render_template('404.html'), 404

if __name__ == "__main__":
    # Get default port from environment variable or use 8080
    default_port = int(os.environ.get('PORT', 8080))
    port = default_port
    max_port_attempts = 10
    
    # Try to find an available port
    for attempt in range(max_port_attempts):
        try:
            logger.info(f"Attempting to start Flask server on port {port}")
            app.run(host='0.0.0.0', port=port)
            break  # If we get here, the app started successfully
        except OSError as e:
            if "Address already in use" in str(e) and attempt < max_port_attempts - 1:
                port = default_port + attempt + 1
                logger.warning(f"Port {port-1} is already in use. Trying port {port}...")
            else:
                logger.error(f"Failed to start server after {attempt+1} attempts: {str(e)}")
                raise