import os
import sqlite3
import hashlib
from datetime import datetime
from flask import Flask, render_template, request, redirect, url_for, flash, jsonify, session
from dotenv import load_dotenv

load_dotenv()
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
if not OPENAI_API_KEY:
    print("WARNING: OPENAI_API_KEY not found. Set it in a .env file or environment variables.")

# Configure OpenAI - using only the legacy client to avoid compatibility issues
client = None
openai = None
USING_NEW_OPENAI = False

# Use the legacy OpenAI module which is working
try:
    import openai
    openai.api_key = OPENAI_API_KEY
    print("Using legacy OpenAI module")
except ImportError:
    print("WARNING: OpenAI module not installed. AI features will be limited.")
    openai = None

DB_PATH = 'app_data.db'

app = Flask(__name__)
app.secret_key = os.getenv('FLASK_SECRET', 'change_this_secret')

# ----------------- Database helpers -----------------
def get_db_connection():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    conn = get_db_connection()
    c = conn.cursor()
    
    # Check if users table exists and has required columns
    c.execute("PRAGMA table_info(users)")
    columns = [column[1] for column in c.fetchall()]
    
    if 'users' not in [table[0] for table in c.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()]:
        # Create users table if it doesn't exist
        c.execute('''
            CREATE TABLE users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT,
                email TEXT UNIQUE,
                password TEXT,
                age INTEGER,
                gender TEXT,
                location TEXT,
                created_at TEXT
            )
        ''')
        print("Created users table with password column")
    elif 'password' not in columns:
        # Add password column if it doesn't exist
        c.execute('ALTER TABLE users ADD COLUMN password TEXT')
        print("Added password column to users table")
    
    if 'location' not in columns:
        # Add location column if it doesn't exist
        c.execute('ALTER TABLE users ADD COLUMN location TEXT')
        print("Added location column to users table")
    else:
        # users table exists with location column
        print("Users table already has location column")
    # doctors table
    c.execute('''
        CREATE TABLE IF NOT EXISTS doctors (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            disease TEXT,
            specialist TEXT,
            doctor_name TEXT,
            contact TEXT,
            location TEXT
        )
    ''')
    
    # Check if chats table exists and has required columns
    c.execute("PRAGMA table_info(chats)")
    chat_columns = [column[1] for column in c.fetchall()]
    
    # Create or update chats table
    if 'chats' not in [table[0] for table in c.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()]:
        # Create chats table if it doesn't exist
        c.execute('''
            CREATE TABLE chats (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER,
                user_input TEXT,
                normalized_disease TEXT,
                specialist TEXT,
                suggested_doctor TEXT,
                suggested_medicines TEXT,
                created_at TEXT,
                FOREIGN KEY(user_id) REFERENCES users(id)
            )
        ''')
        print("Created chats table with suggested_medicines column")
    elif 'suggested_medicines' not in chat_columns:
        # Add suggested_medicines column if it doesn't exist
        c.execute('ALTER TABLE chats ADD COLUMN suggested_medicines TEXT')
        print("Added suggested_medicines column to chats table")
    
    # medicines table
    c.execute('''
        CREATE TABLE IF NOT EXISTS medicines (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            disease TEXT,
            medicine_name TEXT,
            dosage TEXT,
            frequency TEXT,
            description TEXT
        )
    ''')
    conn.commit()
    conn.close()

def seed_doctors():
    conn = get_db_connection()
    c = conn.cursor()
    # check if already seeded
    c.execute('SELECT COUNT(*) FROM doctors')
    if c.fetchone()[0] > 0:
        # Clear existing data to reseed with more doctors
        c.execute('DELETE FROM doctors')
        print("Reseeding doctors table with more data")

    sample = [
        # Diabetes specialists
        ('diabetes', 'Endocrinologist', 'Dr. Ali Khan - City Hospital', '0300-1111111', 'Karachi'),
        ('diabetes', 'Endocrinologist', 'Dr. Sarah Ahmed - Diabetes Care Center', '0300-2222222', 'Karachi'),
        ('diabetes', 'Endocrinologist', 'Dr. Faisal Malik - Medical Complex', '0300-3333333', 'Lahore'),
        ('diabetes', 'Endocrinologist', 'Dr. Nadia Hussain - Health Point', '0300-4444444', 'Islamabad'),
        
        # Heart specialists
        ('heart attack', 'Cardiologist', 'Dr. Imran - Heart Care', '0301-2222222', 'Lahore'),
        ('heart attack', 'Cardiologist', 'Dr. Asma Khan - Cardiac Center', '0301-3333333', 'Karachi'),
        ('heart attack', 'Cardiologist', 'Dr. Tariq Mahmood - City Hospital', '0301-4444444', 'Islamabad'),
        
        # Skin specialists
        ('skin allergy', 'Dermatologist', 'Dr. Ayesha - Skin Glow', '0302-3333333', 'Islamabad'),
        ('skin allergy', 'Dermatologist', 'Dr. Zubair Ahmed - Derma Care', '0302-4444444', 'Karachi'),
        ('skin allergy', 'Dermatologist', 'Dr. Saima Malik - Skin Solutions', '0302-5555555', 'Lahore'),
        
        # Lung specialists
        ('asthma', 'Pulmonologist', 'Dr. Hamid - Lungs Care', '0303-4444444', 'Karachi'),
        ('asthma', 'Pulmonologist', 'Dr. Farah Khan - Respiratory Center', '0303-5555555', 'Lahore'),
        
        # General physicians
        ('fever', 'General Physician', 'Dr. Bilal - Green Clinic', '0304-5555555', 'Karachi'),
        ('fever', 'General Physician', 'Dr. Sana Ali - Family Clinic', '0304-6666666', 'Lahore'),
        ('fever', 'General Physician', 'Dr. Kamran - Health First', '0304-7777777', 'Islamabad')
    ]
    c.executemany('INSERT INTO doctors (disease, specialist, doctor_name, contact, location) VALUES (?, ?, ?, ?, ?)', sample)
    conn.commit()
    conn.close()

def seed_medicines():
    conn = get_db_connection()
    c = conn.cursor()
    # check if already seeded
    c.execute('SELECT COUNT(*) FROM medicines')
    if c.fetchone()[0] > 0:
        # Clear existing data to reseed
        c.execute('DELETE FROM medicines')
        print("Reseeding medicines table with data")

    sample = [
        # Diabetes medicines
        ('diabetes', 'Metformin', '500-1000 mg', 'Twice daily', 'First-line medication for type 2 diabetes. Take with meals to reduce stomach upset.'),
        ('diabetes', 'Insulin Glargine', '10-20 units', 'Once daily', 'Long-acting insulin. Inject at the same time each day.'),
        ('diabetes', 'Glipizide', '5-10 mg', 'Once daily', 'Stimulates insulin release. Take before breakfast.'),
        
        # Heart attack medicines
        ('heart attack', 'Aspirin', '81-325 mg', 'Once daily', 'Blood thinner to prevent clots. Take with food.'),
        ('heart attack', 'Atorvastatin', '10-80 mg', 'Once daily', 'Cholesterol-lowering medication. Take in the evening.'),
        ('heart attack', 'Metoprolol', '25-100 mg', 'Twice daily', 'Beta-blocker to reduce heart workload. Do not stop suddenly.'),
        
        # Skin allergy medicines
        ('skin allergy', 'Cetirizine', '10 mg', 'Once daily', 'Antihistamine for allergies. May cause drowsiness.'),
        ('skin allergy', 'Hydrocortisone Cream', 'Apply thin layer', '2-3 times daily', 'Topical steroid for inflammation. Use only on affected areas.'),
        ('skin allergy', 'Fexofenadine', '180 mg', 'Once daily', 'Non-drowsy antihistamine. Take on an empty stomach.'),
        
        # Asthma medicines
        ('asthma', 'Albuterol Inhaler', '1-2 puffs', 'Every 4-6 hours as needed', 'Quick-relief bronchodilator for asthma attacks.'),
        ('asthma', 'Fluticasone Inhaler', '1-2 puffs', 'Twice daily', 'Preventive steroid inhaler. Rinse mouth after use.'),
        ('asthma', 'Montelukast', '10 mg', 'Once daily at bedtime', 'Leukotriene modifier to prevent asthma symptoms.'),
        
        # Fever medicines
        ('fever', 'Acetaminophen', '500-1000 mg', 'Every 4-6 hours as needed', 'Fever reducer and pain reliever. Do not exceed 4000 mg per day.'),
        ('fever', 'Ibuprofen', '200-400 mg', 'Every 6-8 hours with food', 'Anti-inflammatory and fever reducer. Take with food or milk.'),
        ('fever', 'Naproxen', '220-440 mg', 'Every 8-12 hours', 'Long-acting anti-inflammatory. Take with a full glass of water.')
    ]
    c.executemany('INSERT INTO medicines (disease, medicine_name, dosage, frequency, description) VALUES (?, ?, ?, ?, ?)', sample)
    conn.commit()
    conn.close()

# Initialize DB and seed
init_db()
seed_doctors()
seed_medicines()

# ----------------- Utility: AI normalization -----------------
def keyword_match(user_text, candidates):
    """
    Simple keyword-based disease matching fallback when AI is not available.
    """
    ut = user_text.lower()
    # Exact match
    for d in candidates:
        if d in ut:
            return d
    # Word match
    words = ut.split()
    for d in candidates:
        for w in d.split():
            if w in words:
                return d
    return None

def normalize_disease_with_ai(user_text, candidate_list=None):
    """
    Use OpenAI to map user_text to the closest disease from candidate_list.
    If OPENAI_API_KEY is not set, fallback to simple keyword matching.
    """
    candidates = candidate_list or ['diabetes', 'heart attack', 'skin allergy', 'asthma', 'fever', 'cold']

    # Check if API key is the placeholder value
    if OPENAI_API_KEY == 'your_openai_api_key_here':
        print("WARNING: Using placeholder OpenAI API key. Please set a real key in .env file.")
        return keyword_match(user_text, candidates)
        
    # Simple fallback: keyword search
    if not OPENAI_API_KEY or (not USING_NEW_OPENAI and not openai):
        return keyword_match(user_text, candidates)

    # Build prompt
    prompt = (
        "You are a concise medical assistant.\n"
        "Given a user's free-text description of symptoms or a disease name, map it to exactly one disease from THIS LIST: "
        + ", ".join(candidates)
        + ".\n"
        "Return ONLY the chosen disease from the list (one phrase). If none match, return 'unknown'.\n\n"
        f"User text: \"{user_text}\"\n"
    )

    try:
        if openai:
            try:
                # Check if we're using the older API pattern
                if hasattr(openai, 'ChatCompletion'):
                    resp = openai.ChatCompletion.create(
                        model='gpt-3.5-turbo',
                        messages=[
                            {'role':'system', 'content':'You map user descriptions to a disease name from the provided list.'},
                            {'role':'user', 'content': prompt}
                        ],
                        max_tokens=16,
                        temperature=0
                    )
                    out = resp['choices'][0]['message']['content'].strip().lower()
                else:
                    print("OpenAI module doesn't have ChatCompletion attribute")
                    return keyword_match(user_text, candidates)
            except Exception as e:
                print(f"Error with OpenAI: {e}")
                return keyword_match(user_text, candidates)
        else:
            # If we get here, something went wrong with the OpenAI setup
            print("OpenAI configuration error")
            return keyword_match(user_text, candidates)
            
        # clean output
        if out == 'unknown' or out == 'none':
            return None
        # try to match to candidate if slightly different
        for d in candidates:
            if d in out:
                return d
        return out
    except Exception as e:
        print('OpenAI error:', e)
        return None

# ----------------- Routes -----------------
@app.route('/')
def index():
    # If user is already logged in, redirect to chat
    if 'user_id' in session:
        return redirect(url_for('chat'))
    return render_template('index.html')

@app.route('/about')
def about():
    return render_template('about.html')

def hash_password(password):
    """Hash a password for storing."""
    salt = 'ai_doctor_salt'  # In production, use a proper random salt for each user
    return hashlib.sha256((password + salt).encode()).hexdigest()

@app.route('/register', methods=['POST'])
def register():
    name = request.form.get('name', '').strip()
    email = request.form.get('email', '').strip()
    password = request.form.get('password', '')
    confirm_password = request.form.get('confirm_password', '')
    age = request.form.get('age')
    gender = request.form.get('gender', '').strip()
    location = request.form.get('location', '').strip()

    if not name or not email or not password or not location:
        flash('Name, email, password, and location are required', 'danger')
        return redirect(url_for('index'))
        
    if password != confirm_password:
        flash('Passwords do not match', 'danger')
        return redirect(url_for('index'))

    conn = get_db_connection()
    c = conn.cursor()
    
    # Check if email already exists
    existing_user = c.execute('SELECT id FROM users WHERE email = ?', (email,)).fetchone()
    if existing_user:
        conn.close()
        flash('Email already registered. Please login instead.', 'danger')
        return redirect(url_for('login'))
    
    # Hash the password
    hashed_password = hash_password(password)
    
    # Insert new user
    c.execute('INSERT INTO users (name, email, password, age, gender, location, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)',
              (name, email, hashed_password, age or None, gender, location, datetime.utcnow().isoformat()))
    conn.commit()
    user_id = c.lastrowid
    conn.close()

    # Set session
    session['user_id'] = user_id
    session['user_name'] = name
    session['user_email'] = email
    session['user_location'] = location

    flash('Registration successful! You can start chatting now.', 'success')
    return redirect(url_for('chat'))

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        email = request.form.get('email', '').strip()
        password = request.form.get('password', '')
        
        if not email or not password:
            flash('Email and password are required', 'danger')
            return redirect(url_for('login'))
        
        conn = get_db_connection()
        user = conn.execute('SELECT * FROM users WHERE email = ?', (email,)).fetchone()
        conn.close()
        
        if user and user['password'] == hash_password(password):
            # Set session
            session['user_id'] = user['id']
            session['user_name'] = user['name']
            session['user_email'] = user['email']
            session['user_location'] = user['location']
            
            flash('Login successful!', 'success')
            return redirect(url_for('chat'))
        else:
            flash('Invalid email or password', 'danger')
            return redirect(url_for('login'))
    
    # GET request - show login form
    return render_template('login.html')

@app.route('/logout')
def logout():
    # Clear session
    session.clear()
    flash('You have been logged out', 'success')
    return redirect(url_for('index'))

@app.route('/chat')
def chat():
    # Check if user is logged in
    if 'user_id' not in session:
        flash('Please login to access the chat', 'danger')
        return redirect(url_for('login'))
    
    # Get user data from session
    user_id = session['user_id']
    
    # Get full user data from database
    conn = get_db_connection()
    user = conn.execute('SELECT * FROM users WHERE id = ?', (user_id,)).fetchone()
    conn.close()
    
    return render_template('chat.html', user=user)

@app.route('/history')
def history():
    # Check if user is logged in
    if 'user_id' not in session:
        flash('Please login to view your chat history', 'danger')
        return redirect(url_for('login'))
    
    # Get user data from session
    user_id = session['user_id']
    
    # Get chat history from database
    conn = get_db_connection()
    chats = conn.execute('SELECT * FROM chats WHERE user_id = ? ORDER BY created_at DESC', (user_id,)).fetchall()
    conn.close()
    
    return render_template('history.html', chats=chats)

@app.route('/api/suggest', methods=['POST'])
def api_suggest():
    # Check if user is logged in via session
    if 'user_id' not in session:
        return jsonify({'error':'User not authenticated'}), 401
    
    data = request.json or {}
    user_id = session['user_id']  # Get user_id from session instead of request
    text = data.get('text', '')
    
    print(f"API suggest called with user_id: {user_id}, text: {text}")

    if not text:
        return jsonify({'error':'Empty input'}), 400

    # Get user location from session
    user_location = session.get('user_location')
    
    # If not in session, get from database as fallback
    if not user_location:
        conn = get_db_connection()
        user = conn.execute('SELECT location FROM users WHERE id = ?', (user_id,)).fetchone()
        if user and user['location']:
            user_location = user['location']
            # Update session
            session['user_location'] = user_location

    # fetch candidate diseases from DB
    conn = get_db_connection()
    rows = conn.execute('SELECT DISTINCT disease FROM doctors').fetchall()
    candidates = [r['disease'] for r in rows]

    normalized = normalize_disease_with_ai(text, candidate_list=candidates)

    # if not normalized, simple keyword fallback
    if not normalized:
        normalized = None

    # search doctors based on location and disease
    doctors = []
    specialist = None
    primary_doctor = None
    
    if normalized:
        # First try to find doctors in the user's location
        if user_location:
            rows = conn.execute('SELECT * FROM doctors WHERE lower(disease) = ? AND location = ?', 
                           (normalized.lower(), user_location)).fetchall()
            if rows:
                for row in rows:
                    doctors.append(dict(row))
                specialist = doctors[0]['specialist']
                primary_doctor = doctors[0]['doctor_name']
        
        # If no doctors found in user's location, search without location filter
        if not doctors:
            rows = conn.execute('SELECT * FROM doctors WHERE lower(disease) = ?', (normalized.lower(),)).fetchall()
            if rows:
                for row in rows:
                    doctors.append(dict(row))
                specialist = doctors[0]['specialist']
                primary_doctor = doctors[0]['doctor_name']

    # if nothing found, fallback to general physician suggestion in user's location if possible
    if not doctors:
        if user_location:
            rows = conn.execute("SELECT * FROM doctors WHERE specialist = 'General Physician' AND location = ?", 
                           (user_location,)).fetchall()
            if rows:
                for row in rows:
                    doctors.append(dict(row))
                specialist = doctors[0]['specialist']
                primary_doctor = doctors[0]['doctor_name']
        
        # If still no doctors found, get any general physicians
        if not doctors:
            rows = conn.execute("SELECT * FROM doctors WHERE specialist = 'General Physician'").fetchall()
            if rows:
                for row in rows:
                    doctors.append(dict(row))
                specialist = doctors[0]['specialist']
                primary_doctor = doctors[0]['doctor_name']

    # Get medicine suggestions for the disease
    medicines = []
    if normalized:
        medicine_rows = conn.execute('SELECT * FROM medicines WHERE lower(disease) = ?', (normalized.lower(),)).fetchall()
        if medicine_rows:
            for row in medicine_rows:
                medicines.append(dict(row))
    
    # Format medicines as a string for storage in the database
    suggested_medicines_text = ", ".join([m['medicine_name'] for m in medicines]) if medicines else None
    
    # save chat
    conn.execute('INSERT INTO chats (user_id, user_input, normalized_disease, specialist, suggested_doctor, suggested_medicines, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)',
                 (user_id, text, normalized, specialist, primary_doctor, suggested_medicines_text, datetime.utcnow().isoformat()))
    conn.commit()
    conn.close()

    response = {
        'user_input': text,
        'normalized_disease': normalized,
        'specialist': specialist,
        'doctors': doctors,
        'medicines': medicines,
        'location_match': doctors and user_location and any(d['location'] == user_location for d in doctors)
    }
    return jsonify(response)

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