import socket
import threading
import json
import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox
from datetime import datetime
from tkinter import font as tkfont

class ChatClient:
    def __init__(self, root):
        self.root = root
        self.root.title("Python Chat")
        self.root.geometry("800x600")
        self.root.configure(bg='#f0f0f0')

        self.current_user = None
        self.current_chat = None
        self.is_group_chat = False
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        self.setup_ui()

    def setup_ui(self):
        self.custom_font = tkfont.Font(family="Helvetica", size=10)
        self.bold_font = tkfont.Font(family="Helvetica", size=10, weight="bold")

        self.main_frame = ttk.Frame(self.root)
        self.main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        self.left_panel = ttk.Frame(self.main_frame, width=200)
        self.left_panel.pack(side=tk.LEFT, fill=tk.Y, padx=(0, 10))
        self.left_panel.pack_propagate(False)

        self.right_panel = ttk.Frame(self.main_frame)
        self.right_panel.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)

        self.login_frame = ttk.Frame(self.left_panel)
        self.login_frame.pack(fill=tk.X, pady=(0, 10))

        ttk.Label(self.login_frame, text="Username:").pack(anchor=tk.W)
        self.username_entry = ttk.Entry(self.login_frame)
        self.username_entry.pack(fill=tk.X)

        ttk.Label(self.login_frame, text="Password:").pack(anchor=tk.W, pady=(5, 0))
        self.password_entry = ttk.Entry(self.login_frame, show="*")
        self.password_entry.pack(fill=tk.X)

        self.login_button = ttk.Button(self.login_frame, text="Login", command=self.login)
        self.login_button.pack(fill=tk.X, pady=(5, 0))

        self.register_button = ttk.Button(self.login_frame, text="Register", command=self.register)
        self.register_button.pack(fill=tk.X, pady=(5, 0))

        self.user_list_frame = ttk.LabelFrame(self.left_panel, text="Users")
        self.user_list_frame.pack(fill=tk.BOTH, expand=True)

        self.user_list = tk.Listbox(self.user_list_frame, font=self.custom_font, selectmode=tk.SINGLE)
        self.user_list.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        self.user_list.bind('<<ListboxSelect>>', self.select_user)

        self.group_list_frame = ttk.LabelFrame(self.left_panel, text="Groups")
        self.group_list_frame.pack(fill=tk.BOTH, pady=(10, 0))

        self.group_list = tk.Listbox(self.group_list_frame, font=self.custom_font, selectmode=tk.SINGLE)
        self.group_list.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        self.group_list.bind('<<ListboxSelect>>', self.select_group)

        self.create_group_button = ttk.Button(self.left_panel, text="Create Group", command=self.show_create_group_dialog)
        self.create_group_button.pack(fill=tk.X, pady=(10, 0))

        self.chat_frame = ttk.Frame(self.right_panel)
        self.chat_frame.pack(fill=tk.BOTH, expand=True)

        self.chat_header = ttk.Label(self.chat_frame, text="Select a chat", font=self.bold_font)
        self.chat_header.pack(anchor=tk.W, padx=5, pady=5)

        self.chat_display = scrolledtext.ScrolledText(
            self.chat_frame, wrap=tk.WORD, state=tk.DISABLED,
            font=self.custom_font, bg='white', padx=10, pady=10
        )
        self.chat_display.pack(fill=tk.BOTH, expand=True)

        self.input_frame = ttk.Frame(self.right_panel)
        self.input_frame.pack(fill=tk.X, pady=(10, 0))

        self.message_entry = ttk.Entry(self.input_frame)
        self.message_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 5))
        self.message_entry.bind('<Return>', self.send_message)

        self.send_button = ttk.Button(self.input_frame, text="Send", command=self.send_message)
        self.send_button.pack(side=tk.RIGHT)

        self.update_ui_state(False)

    def update_ui_state(self, logged_in):
        state = tk.NORMAL if logged_in else tk.DISABLED
        self.user_list.config(state=state)
        self.group_list.config(state=state)
        self.message_entry.config(state=state)
        self.send_button.config(state=state)
        self.create_group_button.config(state=state)

        if not logged_in:
            self.user_list.delete(0, tk.END)
            self.group_list.delete(0, tk.END)
            self.chat_header.config(text="Select a chat")
            self.chat_display.config(state=tk.NORMAL)
            self.chat_display.delete(1.0, tk.END)
            self.chat_display.config(state=tk.DISABLED)
            self.current_chat = None

    def login(self):
        username = self.username_entry.get()
        password = self.password_entry.get()

        if not username or not password:
            messagebox.showerror("Error", "Username and password are required")
            return

        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.connect(('localhost', 5555))
            self.socket.send(json.dumps({
                'type': 'login',
                'username': username,
                'password': password
            }).encode('utf-8'))

            response = json.loads(self.socket.recv(1024).decode('utf-8'))
            if response['type'] == 'login_success':
                self.current_user = username
                self.update_ui_state(True)
                threading.Thread(target=self.receive_messages, daemon=True).start()
            else:
                messagebox.showerror("Error", "Login failed - Invalid username or password")
                self.socket.close()
        except Exception as e:
            messagebox.showerror("Error", f"Connection error: {str(e)}")

    def register(self):
        username = self.username_entry.get()
        password = self.password_entry.get()

        if not username or not password:
            messagebox.showerror("Error", "Username and password are required")
            return

        try:
            temp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            temp_socket.connect(('localhost', 5555))
            temp_socket.send(json.dumps({
                'type': 'register',
                'username': username,
                'password': password
            }).encode('utf-8'))

            response = json.loads(temp_socket.recv(1024).decode('utf-8'))
            temp_socket.close()

            if response['type'] == 'register_success':
                messagebox.showinfo("Success", "Registration successful. Please login now.")
            else:
                messagebox.showerror("Error", "Registration failed - username may be taken or invalid")
        except Exception as e:
            messagebox.showerror("Error", f"Connection error: {str(e)}")

    def receive_messages(self):
        while True:
            try:
                message = self.socket.recv(1024).decode('utf-8')
                if not message:
                    break

                data = json.loads(message)
                if data['type'] == 'user_list':
                    self.user_list.delete(0, tk.END)
                    for user in data['users']:
                        if user != self.current_user:
                            self.user_list.insert(tk.END, user)
                elif data['type'] == 'group_list':
                    self.group_list.delete(0, tk.END)
                    for group in data['groups']:
                        self.group_list.insert(tk.END, group)
                elif data['type'] == 'message':
                    self.display_message(
                        data['sender'],
                        data['receiver'],
                        data['message'],
                        data['is_group']
                    )
                elif data['type'] == 'initial_messages':
                    for msg in data['messages']:
                        self.display_message(
                            msg['sender'],
                            msg['receiver'],
                            msg['message'],
                            msg['is_group'],
                            msg['timestamp']
                        )
                elif data['type'] == 'server_shutdown':
                    messagebox.showinfo("Server", "Server is shutting down")
                    self.on_closing()
                    break
            except Exception as e:
                print(f"Error receiving message: {e}")
                break

    def display_message(self, sender, receiver, message, is_group, timestamp=None):
        timestamp = timestamp or datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        if (is_group and receiver == self.current_chat) or \
                (not is_group and ((sender == self.current_chat and receiver == self.current_user) or \
                                   (sender == self.current_user and receiver == self.current_chat))):

            self.chat_display.config(state=tk.NORMAL)
            if sender == self.current_user:
                self.chat_display.tag_config('right', justify='right', foreground='blue')
                self.chat_display.insert(tk.END, f"{message} ({timestamp})\n", 'right')
            else:
                sender_label = f"{sender} (group)" if is_group and sender != self.current_chat else sender
                self.chat_display.tag_config('left', justify='left', foreground='green')
                self.chat_display.insert(tk.END, f"{sender_label}: {message} ({timestamp})\n", 'left')
            self.chat_display.config(state=tk.DISABLED)
            self.chat_display.see(tk.END)

    def select_user(self, event):
        selection = event.widget.curselection()
        if selection:
            self.current_chat = event.widget.get(selection[0])
            self.is_group_chat = False
            self.chat_header.config(text=f"Chat with {self.current_chat}")
            self.chat_display.config(state=tk.NORMAL)
            self.chat_display.delete(1.0, tk.END)
            self.chat_display.config(state=tk.DISABLED)

    def select_group(self, event):
        selection = event.widget.curselection()
        if selection:
            self.current_chat = event.widget.get(selection[0])
            self.is_group_chat = True
            self.chat_header.config(text=f"Group: {self.current_chat}")
            self.chat_display.config(state=tk.NORMAL)
            self.chat_display.delete(1.0, tk.END)
            self.chat_display.config(state=tk.DISABLED)

    def send_message(self, event=None):
        if not self.current_chat or not self.message_entry.get():
            return

        message = self.message_entry.get()
        self.socket.send(json.dumps({
            'type': 'message',
            'sender': self.current_user,
            'receiver': self.current_chat,
            'message': message,
            'is_group': self.is_group_chat
        }).encode('utf-8'))

        self.display_message(
            self.current_user,
            self.current_chat,
            message,
            self.is_group_chat
        )
        self.message_entry.delete(0, tk.END)

    def show_create_group_dialog(self):
        if not self.current_user:
            return

        dialog = tk.Toplevel(self.root)
        dialog.title("Create Group")
        dialog.geometry("300x300")

        ttk.Label(dialog, text="Group Name:").pack(anchor=tk.W, padx=10, pady=(10, 0))
        group_name_entry = ttk.Entry(dialog)
        group_name_entry.pack(fill=tk.X, padx=10, pady=(0, 10))

        ttk.Label(dialog, text="Select Members:").pack(anchor=tk.W, padx=10, pady=(10, 0))

        members_frame = ttk.Frame(dialog)
        members_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=(0, 10))

        members_list = tk.Listbox(members_frame, selectmode=tk.MULTIPLE)
        members_list.pack(fill=tk.BOTH, expand=True)

        for i in range(self.user_list.size()):
            members_list.insert(tk.END, self.user_list.get(i))

        def create_group():
            group_name = group_name_entry.get()
            selected_indices = members_list.curselection()
            selected_members = [members_list.get(i) for i in selected_indices]

            if not group_name or not selected_members:
                messagebox.showerror("Error", "Group name and at least one member are required")
                return

            self.socket.send(json.dumps({
                'type': 'create_group',
                'group_name': group_name,
                'members': selected_members + [self.current_user]
            }).encode('utf-8'))
            dialog.destroy()

        button_frame = ttk.Frame(dialog)
        button_frame.pack(fill=tk.X, padx=10, pady=(0, 10))

        ttk.Button(button_frame, text="Create", command=create_group).pack(side=tk.RIGHT)
        ttk.Button(button_frame, text="Cancel", command=dialog.destroy).pack(side=tk.RIGHT, padx=(0, 5))

    def on_closing(self):
        if hasattr(self, 'socket') and self.current_user:
            try:
                self.socket.close()
            except:
                pass
        self.root.destroy()

if __name__ == "__main__":
    root = tk.Tk()
    client = ChatClient(root)
    root.protocol("WM_DELETE_WINDOW", client.on_closing)
    root.mainloop()