from dataclasses import dataclass, field
from uuid import UUID, uuid4
from src.customer import Customer
from src.models import NKodePolicy, KeypadSize
from src.user import User
from src.user_cipher import UserCipher
from src.user_signup_session import UserSignupSession
from src.user_keypad import UserKeypad
from src.customer_cipher import CustomerCipher
import numpy as np


@dataclass
class NKodeAPI:
    customers: dict[UUID, Customer] = field(default_factory=dict)
    signup_sessions: dict[UUID, UserSignupSession] = field(default_factory=dict)

    def create_new_customer(self, keypad_size: KeypadSize, nkode_policy: NKodePolicy) -> UUID:
        new_customer = Customer.create(
            cipher=CustomerCipher.create(keypad_size),
            nkode_policy=nkode_policy
        )
        self.customers[new_customer.customer_id] = new_customer
        return new_customer.customer_id

    def generate_signup_keypad(self, customer_id: UUID) -> tuple[UUID, np.ndarray]:
        if customer_id not in self.customers.keys():
            raise ValueError(f"Customer with ID '{customer_id}' does not exist")
        customer = self.customers[customer_id]
        login_keypad = UserKeypad.create(customer.cipher.keypad_size)
        set_keypad = login_keypad.sign_up_keypad()
        new_session = UserSignupSession(
            session_id=uuid4(),
            login_keypad=login_keypad,
            set_keypad=set_keypad.keypad,
            customer_id=customer_id,
            keypad_size=set_keypad.keypad_size,
        )
        self.signup_sessions[new_session.session_id] = new_session
        return new_session.session_id, new_session.set_keypad

    def set_nkode(
            self,
            username: str,
            customer_id: UUID,
            key_selection: list[int],
            session_id: UUID
    ) -> np.ndarray:
        if customer_id not in self.customers.keys():
            raise ValueError(f"Customer ID {customer_id} not found")
        customer = self.customers[customer_id]
        if username in customer.users.keys():
            raise ValueError(f"Username '{username}' already exists for this customer")
        if session_id not in self.signup_sessions.keys():
            raise ValueError(f"Session ID {session_id} not found")
        self.signup_sessions[session_id].set_user_nkode(username, key_selection)
        return self.signup_sessions[session_id].confirm_keypad

    def confirm_nkode(
            self,
            username: str,
            customer_id: UUID,
            confirm_key_entry: list[int],
            session_id: UUID
    ) -> bool:
        if session_id not in self.signup_sessions.keys():
            raise AssertionError(f"Session ID {session_id} not found in signup sessions")
        session = self.signup_sessions[session_id]
        if customer_id != session.customer_id:
            raise AssertionError(f"Customer ID mismatch: {customer_id} vs {session.customer_id}")
        if username != session.username:
            raise AssertionError(f"Username mismatch: {username} vs {session.username}")
        customer = self.customers[customer_id]
        passcode = self.signup_sessions[session_id].deduce_passcode(confirm_key_entry)
        new_user_keys = UserCipher.create(
            customer.cipher.keypad_size,
            customer.cipher.set_key,
            customer.nkode_policy.max_nkode_len
        )
        enciphered_passcode = new_user_keys.encipher_nkode(passcode, customer.cipher)
        new_user = User(
            username=username,
            enciphered_passcode=enciphered_passcode,
            cipher=new_user_keys,
            user_keypad=self.signup_sessions[session_id].login_keypad,
        )
        self.customers[customer_id].add_new_user(new_user)
        del self.signup_sessions[session_id]
        return True

    def get_login_keypad(self, username: str, customer_id: UUID) -> np.ndarray:
        if customer_id not in self.customers.keys():
            raise ValueError("Customer ID not found")
        customer = self.customers[customer_id]
        if username not in customer.users.keys():
            raise ValueError("Username not found")
        user = customer.users[username]
        # TODO: implement split_keypad_shuffle()
        return user.user_keypad.keypad

    def login(self, customer_id: UUID, username: str, key_selection: list[int]) -> bool:
        if customer_id not in self.customers.keys():
            raise ValueError("Customer ID not found")
        customer = self.customers[customer_id]
        return customer.valid_key_entry(username, key_selection)

    def renew_keys(self, customer_id: UUID) -> bool:
        if customer_id not in self.customers.keys():
            raise ValueError("Customer ID not found")
        return self.customers[customer_id].renew_keys()