import socket
import threading
import sys
import os
import shutil
from datetime import datetime

# Colorama for Windows color support
try:
    from colorama import init, Fore, Back, Style
    init(autoreset=True)
    HAS_COLOR = True
except ImportError:
    HAS_COLOR = False

SERVER_HOST = "YOUR_SERVER_IP"   # <-- Replace with your Railway host
SERVER_PORT = 9999                # <-- Replace with your Railway port

# ─── Color helpers ────────────────────────────────────────────────────────────

def c(text, *codes):
    if not HAS_COLOR:
        return text
    return "".join(codes) + text + Style.RESET_ALL

def term_width():
    return shutil.get_terminal_size((80, 24)).columns

# ─── UI Drawing ───────────────────────────────────────────────────────────────

def clear_screen():
    os.system("cls" if os.name == "nt" else "clear")

def draw_banner():
    clear_screen()
    w = term_width()
    bar = "═" * w
    title = "NormansChat"
    subtitle = "Terminal Messenger"
    print(c(bar, Fore.GREEN))
    print(c(title.center(w), Fore.GREEN, Style.BRIGHT))
    print(c(subtitle.center(w), Fore.WHITE))
    print(c(bar, Fore.GREEN))
    print()

def draw_chat_header(partner_name):
    w = term_width()
    bar = "─" * w
    print(c(bar, Fore.GREEN))
    online = "● ONLINE"
    left = f"  💬 {partner_name}"
    right = f"{online}  "
    gap = w - len(left) - len(right)
    print(c(left, Fore.WHITE, Style.BRIGHT) + " " * max(gap, 1) + c(right, Fore.GREEN))
    print(c(bar, Fore.GREEN))
    print(c("  Type /quit to leave\n", Fore.WHITE))

def timestamp():
    return datetime.now().strftime("%H:%M")

def print_my_message(msg):
    """Right-aligned green bubble — like WhatsApp sent messages."""
    w = term_width()
    ts = timestamp()
    bubble = f"  {msg}  {ts} ✓"
    padding = w - len(bubble) - 2
    line = " " * max(padding, 0) + c(f"  {msg}  ", Fore.BLACK, Back.GREEN) + c(f"{ts} ✓  ", Fore.GREEN)
    sys.stdout.write(f"\r\033[K{line}\n")
    sys.stdout.flush()

def print_their_message(name, msg):
    """Left-aligned white bubble — like WhatsApp received messages."""
    ts = timestamp()
    name_line = c(f"  {name}", Fore.GREEN, Style.BRIGHT)
    bubble = c(f"  {msg}  ", Fore.BLACK, Back.WHITE) + c(f"  {ts}", Fore.WHITE)
    sys.stdout.write(f"\r\033[K{name_line}\n{bubble}\n")
    sys.stdout.flush()

def print_system(msg):
    """Centered gray system message — like WhatsApp's 'encryption' notices."""
    w = term_width()
    text = f"  {msg}  "
    padding = (w - len(text)) // 2
    line = " " * max(padding, 0) + c(text, Fore.YELLOW)
    sys.stdout.write(f"\r\033[K{line}\n")
    sys.stdout.flush()

def show_input_prompt(name):
    prompt = c(f"\n  ✏  {name}: ", Fore.GREEN, Style.BRIGHT)
    sys.stdout.write(prompt)
    sys.stdout.flush()

# ─── Networking ───────────────────────────────────────────────────────────────

# Shared state
my_name = ""
partner_name = ""
chat_started = False
input_lock = threading.Lock()

def receive_messages(sock, stop_event):
    global my_name, partner_name, chat_started
    buffer = ""
    try:
        while not stop_event.is_set():
            try:
                sock.settimeout(1.0)
                data = sock.recv(4096)
            except socket.timeout:
                continue
            except:
                break

            if not data:
                break

            text = data.decode(errors="ignore")
            buffer += text

            # Process line by line
            while "\n" in buffer:
                line, buffer = buffer.split("\n", 1)
                line = line.strip().rstrip("\r").strip()
                if not line:
                    continue

                # Detect partner messages (format: "Name: message")
                if chat_started and ": " in line and not line.startswith("[") and not line.startswith("You"):
                    # Strip any leading whitespace/formatting
                    clean = line.strip()
                    if ": " in clean:
                        sender, msg = clean.split(": ", 1)
                        sender = sender.strip()
                        msg = msg.strip()
                        if sender and msg and sender != my_name:
                            with input_lock:
                                sys.stdout.write("\r\033[K")  # clear input line
                                print_their_message(sender, msg)
                                show_input_prompt(my_name)
                            continue

                # System / server messages
                clean = line.strip()
                if not clean:
                    continue

                # Detect chat start
                if "joined the room" in clean or "Chatting with" in clean or "connected" in clean.lower():
                    chat_started = True
                    # Extract partner name if possible
                    if "joined the room" in clean:
                        part = clean.replace("[", "").replace("]", "").replace("joined the room!", "").strip()
                        if part:
                            partner_name = part
                    elif "Chatting with" in clean:
                        part = clean.split("Chatting with")[-1].strip().rstrip(".")
                        if part:
                            partner_name = part

                # Skip "You: " prompts — we handle input ourselves
                if clean in ("You:", "You: ") or clean.startswith("You: ") and len(clean) < 7:
                    continue

                with input_lock:
                    if chat_started:
                        sys.stdout.write("\r\033[K")
                    print_system(clean)
                    if chat_started:
                        show_input_prompt(my_name)

    except:
        pass
    finally:
        stop_event.set()


def send_messages(sock, stop_event):
    global chat_started
    try:
        while not stop_event.is_set():
            try:
                if chat_started:
                    show_input_prompt(my_name)
                msg = input("")
            except (EOFError, KeyboardInterrupt):
                break

            if stop_event.is_set():
                break

            msg = msg.strip()
            if not msg:
                continue

            try:
                sock.send((msg + "\r\n").encode())
            except:
                break

            if msg.lower() == "/quit":
                break

            # Print our own message right-aligned
            if chat_started:
                with input_lock:
                    print_my_message(msg)

    except:
        pass
    finally:
        stop_event.set()


# ─── Setup flow (name + room ID) ──────────────────────────────────────────────

def setup_flow(sock):
    """Handle the name + room ID prompts before chat starts."""
    global my_name
    buffer = ""
    step = 0  # 0=waiting for name prompt, 1=waiting for room prompt, 2=done

    while step < 2:
        try:
            sock.settimeout(10.0)
            data = sock.recv(4096)
        except:
            return False

        if not data:
            return False

        text = data.decode(errors="ignore")
        buffer += text
        sys.stdout.write(text)
        sys.stdout.flush()

        if step == 0 and "display name" in buffer.lower():
            name = input("")
            my_name = name.strip()
            sock.send((name + "\r\n").encode())
            step = 1

        elif step == 1 and "room id" in buffer.lower():
            room = input("")
            sock.send((room + "\r\n").encode())
            step = 2

    return True


# ─── Main ─────────────────────────────────────────────────────────────────────

def main():
    global my_name, chat_started

    if os.name == "nt":
        os.system("")  # enable ANSI on Windows

    if not HAS_COLOR:
        print("  Tip: Run 'pip install colorama' for colors!\n")

    draw_banner()
    print(c("  Connecting to server...\n", Fore.WHITE))

    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect((SERVER_HOST, SERVER_PORT))
    except ConnectionRefusedError:
        print(c(f"\n  [ERROR] Could not connect to {SERVER_HOST}:{SERVER_PORT}", Fore.RED))
        print(c("  Make sure the server is running.", Fore.WHITE))
        input("\n  Press Enter to exit...")
        sys.exit(1)
    except socket.gaierror:
        print(c(f"\n  [ERROR] Could not resolve: {SERVER_HOST}", Fore.RED))
        input("\n  Press Enter to exit...")
        sys.exit(1)

    # Handle name + room setup interactively
    if not setup_flow(sock):
        print(c("\n  Connection lost during setup.", Fore.RED))
        sys.exit(1)

    # Wait a moment for server to confirm connection
    import time
    time.sleep(1.5)

    # Now draw the chat UI
    clear_screen()
    draw_chat_header(partner_name if partner_name else "Partner")

    stop_event = threading.Event()
    recv_thread = threading.Thread(target=receive_messages, args=(sock, stop_event), daemon=True)
    send_thread = threading.Thread(target=send_messages, args=(sock, stop_event), daemon=True)

    recv_thread.start()
    send_thread.start()

    try:
        stop_event.wait()
    except KeyboardInterrupt:
        pass

    print(c("\n\n  Goodbye! 👋\n", Fore.GREEN, Style.BRIGHT))
    try:
        sock.close()
    except:
        pass


if __name__ == "__main__":
    main()
