import sqlite3
from flask import Flask, render_template, jsonify, request, redirect, url_for
from collections import defaultdict
import os
import json
import re
import wordninja
from typing import List, Set
from flask_login import LoginManager, UserMixin, login_user, logout_user, login_required, current_user
from werkzeug.security import generate_password_hash, check_password_hash

class CustomWordNinja:
    def __init__(self, custom_words: List[str] = None):
        self.custom_words = []
        if custom_words:
            # Store custom words with original capitalization, sorted by length
            self.custom_words = sorted(custom_words, key=len, reverse=True)
    
    def split(self, text: str) -> str:
        working_text = text
        
        # First handle exact custom words to preserve capitalization
        for word in self.custom_words:
            pattern = re.compile(word, re.IGNORECASE)
            working_text = pattern.sub(f' {word} ', working_text)
        
        # Clean up spaces
        working_text = ' '.join(working_text.split())
        
        # For remaining text, use wordninja
        parts = []
        for part in working_text.split():
            if part in self.custom_words:
                parts.append(part)
            else:
                split_parts = wordninja.split(part)
                parts.extend(split_parts)
        
        return ' '.join(parts)

POKEMON_PROPER_NOUNS = {
    "Alola",
    "Kanto",
    "Johto",
    "Hoenn",
    "Sinnoh",
    "Unova",
    "Kalos",
    "Galar",
    "Hisui",
    "Paldea",
    "Augurite",
    "Electirizer",
    "Magmarizer"
    # Add any other proper nouns that should be preserved
}

app = Flask(__name__)
app.secret_key = '8PQ46ThghoOUsmr4iPMHvrT91FESh9kHu'  # Change this to a secure secret key
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = 'login'

class User(UserMixin):
    def __init__(self, id, username):
        self.id = id
        self.username = username

@login_manager.user_loader
def load_user(user_id):
    conn = sqlite3.connect('user_data.db')
    cursor = conn.cursor()
    cursor.execute('SELECT id, username FROM users WHERE id = ?', (user_id,))
    user = cursor.fetchone()
    conn.close()
    
    if user:
        return User(user[0], user[1])
    return None

# Add login/register routes
@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        
        conn = sqlite3.connect('user_data.db')
        cursor = conn.cursor()
        cursor.execute('SELECT id, password_hash FROM users WHERE username = ?', (username,))
        user = cursor.fetchone()
        conn.close()
        
        if user and check_password_hash(user[1], password):
            login_user(User(user[0], username))
            return redirect(url_for('index'))
            
        return 'Invalid username or password'
        
    return render_template('login.html')

@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        
        conn = sqlite3.connect('user_data.db')
        cursor = conn.cursor()
        
        try:
            cursor.execute(
                'INSERT INTO users (username, password_hash) VALUES (?, ?)',
                (username, generate_password_hash(password))
            )
            conn.commit()
            conn.close()
            return redirect(url_for('login'))
        except sqlite3.IntegrityError:
            conn.close()
            return 'Username already exists'
            
    return render_template('register.html')

@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('index'))

def load_pokemon_data():
    pokemon_list = []
    
    conn = sqlite3.connect('pokemon_forms.db')
    cursor = conn.cursor()

    # First query: Get all Pokémon data
    cursor.execute('''
        SELECT 
            pf.national_dex, pf.name, pf.form_name, pf.PFIC, pf.generation,
            ps.storable_in_home
        FROM pokemon_forms pf
        JOIN pokemon_storage ps ON pf.PFIC = ps.PFIC

        WHERE ps.storable_in_home = 1
        ORDER BY pf.PFIC
    ''')
    
    pokemon_data = cursor.fetchall()

    # Process the data
    current_group = []
    current_generation = None
    current_dex_number = None
    pokemon_forms = []
    
    for row in pokemon_data:
        national_dex, name, form_name, pfic, generation, storable_in_home = row

        mark_id = None
        mark_icon = None
        
        cursor.execute('''
            SELECT mark_id
            FROM form_marks
            WHERE pfic = ?
        ''', (pfic,))
        result = cursor.fetchone()
        mark_id = result[0] if result else None

        if mark_id:
            cursor.execute('''
                SELECT icon_path
                FROM marks
                WHERE id = ?
            ''', (mark_id,))
            result = cursor.fetchone()
            mark_icon = result[0] if result else None

        pokemon = {
            'pfic': pfic,
            'ID': national_dex,
            'Name': name,
            'Form': form_name if form_name else "Default",
            'Image': f"images/pokemon/{pfic}.png",
            'IsDefault': form_name is None or form_name == '',
            'Generation': generation,
            'StorableInHome': storable_in_home,
            'MarkIcon': mark_icon,
            'MarkID': mark_id
        }
        
        # Add the Pokémon to the current group
        if national_dex != current_dex_number:
            if pokemon_forms:
                for form in pokemon_forms:
                    current_group.append(form)
                    if len(current_group) == 30:
                        pokemon_list.append(current_group)
                        current_group = []
            pokemon_forms = []
            current_dex_number = national_dex
            
            if form_name is None or form_name == '':
                if current_generation is None or generation != current_generation:
                    if current_group:
                        while len(current_group) < 30:
                            current_group.append(None)  # Add empty slots
                        pokemon_list.append(current_group)
                    current_group = []
                    current_generation = generation
        
        pokemon_forms.append(pokemon)
    
    # Add the last set of forms
    for form in pokemon_forms:
        current_group.append(form)
        if len(current_group) == 30:
            pokemon_list.append(current_group)
            current_group = []
    
    # Add any remaining Pokémon
    if current_group:
        while len(current_group) < 30:
            current_group.append(None)  # Add empty slots to the last group
        pokemon_list.append(current_group)

    conn.close()
    return pokemon_list

@app.route('/')
def index():
    pokemon_list = load_pokemon_data()

    splitter = CustomWordNinja(POKEMON_PROPER_NOUNS)
    
    try:
        with open('efficiency_plan.json', 'r', encoding='utf-8') as f:
            efficiency_plan = json.load(f)
            
        conn = sqlite3.connect('pokemon_forms.db')
        cursor = conn.cursor()
        
        def get_evolution_methods(from_pfic, to_pfic):
            # Get direct evolution method
            cursor.execute('''
                SELECT method, to_pfic 
                FROM evolution_chains 
                WHERE from_pfic = ? AND to_pfic = ?
            ''', (from_pfic, to_pfic))
            direct = cursor.fetchone()
            
            if direct:
                words = splitter.split(direct[0])
                return [words]
                
            # Try to find indirect evolution path
            cursor.execute('''
                WITH RECURSIVE evolution_path AS (
                    -- Base case: direct evolutions from start
                    SELECT from_pfic, to_pfic, method, 1 as depth
                    FROM evolution_chains
                    WHERE from_pfic = ?
                    
                    UNION ALL
                    
                    -- Recursive case: follow chain
                    SELECT e.from_pfic, e.to_pfic, e.method, ep.depth + 1
                    FROM evolution_chains e
                    JOIN evolution_path ep ON e.from_pfic = ep.to_pfic
                    WHERE ep.depth < 3  -- Prevent infinite loops
                )
                SELECT method
                FROM evolution_path
                WHERE to_pfic = ?
                ORDER BY depth;
            ''', (from_pfic, to_pfic))
            
            methods = cursor.fetchall()
            if methods:
                # Clean up each method string
                cleaned_methods = []
                for method in methods:
                    # Split and rejoin with spaces
                    words = splitter.split(method[0])
                    cleaned_methods.append(words)
                return cleaned_methods
            return ['Evolution']
        
        # Enhance the plan with evolution methods
        for game in efficiency_plan:
            for pokemon in game['pokemon']:
                if pokemon.get('evolve_to'):
                    for evolution in pokemon['evolve_to']:
                        methods = get_evolution_methods(pokemon['pfic'], evolution['pfic'])
                        evolution['method'] = ' → '.join(methods)
        
        conn.close()
            
    except FileNotFoundError:
        efficiency_plan = []
    
    return render_template('index.html', 
                         grouped_pokemon=pokemon_list, 
                         efficiency_plan=efficiency_plan)

@app.route('/pokemon/<string:pfic>')
def pokemon_details(pfic):
    conn = sqlite3.connect('pokemon_forms.db')
    cursor = conn.cursor()

    cursor.execute('''
        SELECT pf.form_name, g.name, g.id as game_id, e.location, e.day, e.time, 
               e.dual_slot, e.static_encounter, e.static_encounter_count, 
               e.extra_text, e.stars, e.fishing, e.rods, e.starter
        FROM pokemon_forms pf
        JOIN encounters e ON pf.PFIC = e.pfic
        JOIN games g ON e.game_id = g.id
        WHERE pf.pfic = ?
        ORDER BY g.id, pf.form_name, e.location
    ''', (pfic,))

    encounters = [
        {
            'form': form,
            'game': game,
            'game_id': game_id,
            'location': location,
            'day': day,
            'time': time,
            'dual_slot': dual_slot,
            'static_encounter': static_encounter,
            'static_encounter_count': static_encounter_count,
            'extra_text': extra_text,
            'stars': stars,
            'fishing': fishing,
            'rods': rods,
            'starter': starter
        }
        for form, game, game_id, location, day, time, dual_slot, static_encounter, static_encounter_count, 
            extra_text, stars, fishing, rods, starter in cursor.fetchall()
    ]

    conn.close()
    return jsonify(encounters)

def init_user_db():
    conn = sqlite3.connect('user_data.db')
    cursor = conn.cursor()
    
    # Create users table
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # Modify caught_pokemon table to include user_id
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS caught_pokemon (
            user_id INTEGER,
            pfic TEXT,
            caught_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            PRIMARY KEY (user_id, pfic),
            FOREIGN KEY (user_id) REFERENCES users(id)
        )
    ''')
    
    conn.commit()
    conn.close()

@app.route('/toggle_catch/<string:pfic>', methods=['POST'])
@login_required
def toggle_catch(pfic):
    # Verify PFIC exists
    ref_conn = sqlite3.connect('pokemon_forms.db')
    ref_cursor = ref_conn.cursor()
    
    ref_cursor.execute('SELECT 1 FROM pokemon_forms WHERE pfic = ?', (pfic,))
    if not ref_cursor.fetchone():
        ref_conn.close()
        return jsonify({'error': 'Invalid PFIC'}), 400
    ref_conn.close()
    
    user_conn = sqlite3.connect('user_data.db')
    user_cursor = user_conn.cursor()
    
    try:
        user_cursor.execute(
            'SELECT 1 FROM caught_pokemon WHERE user_id = ? AND pfic = ?', 
            (current_user.id, pfic)
        )
        result = user_cursor.fetchone()
        
        if result:
            user_cursor.execute(
                'DELETE FROM caught_pokemon WHERE user_id = ? AND pfic = ?', 
                (current_user.id, pfic)
            )
            status = 'uncaught'
        else:
            user_cursor.execute(
                'INSERT INTO caught_pokemon (user_id, pfic) VALUES (?, ?)', 
                (current_user.id, pfic)
            )
            status = 'caught'
            
        user_conn.commit()
        return jsonify({'status': status, 'pfic': pfic})
        
    except Exception as e:
        user_conn.rollback()
        return jsonify({'error': str(e)}), 500
    finally:
        user_conn.close()

@app.route('/get_caught_pokemon')
@login_required
def get_caught_pokemon():
    user_conn = sqlite3.connect('user_data.db')
    user_cursor = user_conn.cursor()
    
    user_cursor.execute(
        'SELECT pfic FROM caught_pokemon WHERE user_id = ?', 
        (current_user.id,)
    )
    caught = [row[0] for row in user_cursor.fetchall()]
    
    user_conn.close()
    return jsonify(caught)

if __name__ == '__main__':
    init_user_db()  # Initialize user database on startup
    extra_files = ['.']
    extra_dirs = ['./templates/', './static/']
    for extra_dir in extra_dirs:
        for dirname, dirs, files in os.walk(extra_dir):
            for filename in files:
                filename = os.path.join(dirname, filename)
                if os.path.isfile(filename):
                    extra_files.append(filename)

    app.run(debug=True, extra_files=extra_files)