import asyncio
import logging
import sys
import os

# Add the lib directory to sys.path so we can import the openttd package
sys.path.append(os.path.join(os.path.dirname(__file__), 'lib'))

from openttd import OpenTTDClient
from openttd.protocol import ModifyTimetableFlags

# Configuration
SERVER_HOST = "127.0.0.1"
SERVER_PORT = 3979
SERVER_PASSWORD = "asd"

# Company configuration
COMPANY_ID = 0 # "Én transport"
COMPANY_PASSWORD = "asd123"

# A vehicle owned by COMPANY_ID, used to demonstrate timetable get/set below. Set to a real
# vehicle id to see it in action; leave as None to skip the demonstration.
DEMO_VEHICLE_ID = 7

async def demo_timetable_workflow(client, vehicle_id):
    """A deliberately thorough walk-through of the timetable API: every ModifyTimetableFlags
    variant, clear_field, autofill, timetable start, and lateness reset, on a vehicle assumed
    to have at least two orders (positions 0 and 1)."""
    ORDER_A, ORDER_B = 0, 1

    def show(label):
        print(f"    -> [{label}] {client.get_vehicle_timetable(vehicle_id)}")

    print(f"=== Timetable demo starting for vehicle {vehicle_id} ===")

    # 1. Wait/travel times, in game ticks.
    print("--- Step 1: set wait/travel times ---")
    await client.change_timetable(vehicle_id, ORDER_A, ModifyTimetableFlags.WaitTime, 90)
    await client.change_timetable(vehicle_id, ORDER_A, ModifyTimetableFlags.TravelTime, 240)
    await client.change_timetable(vehicle_id, ORDER_B, ModifyTimetableFlags.WaitTime, 45)
    await asyncio.sleep(0.5)
    show("wait/travel times set")

    # 2. Lock order A's wait time so autofill won't overwrite it later.
    print("--- Step 2: fix order A's wait time ---")
    await client.change_timetable(vehicle_id, ORDER_A, ModifyTimetableFlags.SetWaitFixed, 1)
    await asyncio.sleep(0.5)
    show("order A wait time fixed")

    # 3. Cap order B's speed, then remove the cap again (0 = uncapped).
    print("--- Step 3: cap and uncap order B's speed ---")
    await client.change_timetable(vehicle_id, ORDER_B, ModifyTimetableFlags.TravelSpeed, 80)
    await asyncio.sleep(0.5)
    show("order B speed capped at 80")
    await client.change_timetable(vehicle_id, ORDER_B, ModifyTimetableFlags.TravelSpeed, 0)
    await asyncio.sleep(0.5)
    show("order B speed cap removed")

    # 4. Let the vehicle leave order B early once any cargo is fully loaded.
    print("--- Step 4: change order B's leave type ---")
    await client.change_timetable(vehicle_id, ORDER_B, ModifyTimetableFlags.SetLeaveType, 2)
    await asyncio.sleep(0.5)
    show("order B leave type: leave early if any cargo full")

    # 5. Assign order A to scheduled-dispatch schedule 0, then unassign it again.
    print("--- Step 5: assign and unassign a dispatch schedule ---")
    await client.change_timetable(vehicle_id, ORDER_A, ModifyTimetableFlags.AssignSchedule, 0)
    await asyncio.sleep(0.5)
    show("order A assigned to dispatch schedule 0")
    await client.change_timetable(vehicle_id, ORDER_A, ModifyTimetableFlags.AssignSchedule, 0xFFFFFFFF)
    await asyncio.sleep(0.5)
    show("order A unassigned from dispatch schedule")

    # 6. Clear order B's wait time entirely (distinct from setting it to 0).
    print("--- Step 6: clear order B's wait time ---")
    await client.change_timetable(vehicle_id, ORDER_B, ModifyTimetableFlags.WaitTime, 0, clear_field=True)
    await asyncio.sleep(0.5)
    show("order B wait time cleared")

    # 7. Autofill: start it preserving existing (fixed) wait times, then turn it off again.
    print("--- Step 7: toggle autofill ---")
    await client.autofill_timetable(vehicle_id, autofill=True, preserve_wait_time=True)
    await asyncio.sleep(0.5)
    show("autofill enabled (preserving wait times)")
    await client.autofill_timetable(vehicle_id, autofill=False)
    await asyncio.sleep(0.5)
    show("autofill disabled")

    # 8. Start the timetable for this vehicle only, then restart it for the whole group.
    print("--- Step 8: set timetable start ---")
    await client.set_timetable_start(vehicle_id, timetable_all=False, start_date=1_000_000)
    await asyncio.sleep(0.5)
    show("timetable started (this vehicle only)")
    await client.set_timetable_start(vehicle_id, timetable_all=True, start_date=1_500_000)
    await asyncio.sleep(0.5)
    show("timetable restarted (whole group)")

    # 9. Reset lateness for this vehicle, then for the whole group sharing its orders.
    print("--- Step 9: reset lateness ---")
    await client.set_vehicle_on_time(vehicle_id, apply_to_group=False)
    await asyncio.sleep(0.5)
    show("lateness reset (this vehicle only)")
    await client.set_vehicle_on_time(vehicle_id, apply_to_group=True)
    await asyncio.sleep(0.5)
    show("lateness reset (whole group)")

    print(f"=== Timetable demo finished. Final state for vehicle {vehicle_id}: ===")
    print(f"    {client.get_vehicle_timetable(vehicle_id)}")

async def run_client():
    # 1. Initialize high-level client
    username = sys.argv[1] if len(sys.argv) > 1 else "Modular_Joiner"
    client = OpenTTDClient(host=SERVER_HOST, port=SERVER_PORT, username=username)

    # 2. Setup chat callback (optional)
    def chat_logger(cid, msg):
        print(f">>> [CHAT] <{cid}> {msg}")
    client.on_chat = chat_logger

    try:
        # 3. Connect and initiate handshake
        # The client will handle PAKE auth and encryption automatically
        await client.connect(server_password=SERVER_PASSWORD)
        
        # 4. Configure company join
        # This will happen automatically once the handshake is done
        await client.join_company(company_id=COMPANY_ID, company_password=COMPANY_PASSWORD)

        # 5. Wait for the client to be fully synced (map downloaded, states progressed)
        print(f"--- Joining as {username}... ---")
        await client.joined.wait()
        print(f"--- Successfully joined! Client ID: {client.client_id} ---")

        # 6. Timetable demonstration (requires DEMO_VEHICLE_ID to be owned by COMPANY_ID)
        if DEMO_VEHICLE_ID is not None:
            await demo_timetable_workflow(client, DEMO_VEHICLE_ID)

        # 7. Lifecycle management
        # We wait for either a manual shutdown signal or a 10s timeout
        try:
            await asyncio.wait_for(client.shutdown_event.wait(), timeout=10.0)
        except asyncio.TimeoutError:
            print("--- Finished 10s stay, exiting gracefully ---")
            await client.quit()

    except Exception as e:
        print(f"!!! Error: {e}")

if __name__ == "__main__":
    # Setup global logging
    logging.basicConfig(level=logging.INFO, format='%(levelname)s:%(name)s:%(message)s')
    
    # Run the async loop
    try:
        asyncio.run(run_client())
    except KeyboardInterrupt:
        pass