From 2e831adf3aab5c5934392e90b0f6af0e3edd1c5c Mon Sep 17 00:00:00 2001
From: Isaak <isaak.buslovich@protonmail.ch>
Date: Mon, 20 Nov 2023 00:14:09 +0100
Subject: [PATCH] Still some really weird behaviour, but stock trading is
 working a bit better now than before.

---
 main.py | 75 ++++++++++++++++++++++++++++++---------------------------
 1 file changed, 40 insertions(+), 35 deletions(-)

diff --git a/main.py b/main.py
index c0b7a75..1879932 100644
--- a/main.py
+++ b/main.py
@@ -470,50 +470,55 @@ class Company:
             return "Invalid stock action."
 
     def _buy_stock(self, company_id, amount, total_value, is_ai):
-        """Handles the buying of stocks for the specified company."""
-        available_shares = self._calculate_available_shares()
-        if amount > available_shares:
-            return f"Not enough available shares to buy. Available: {available_shares}"
+        # Calculate the total number of shares currently owned
+        total_shares_owned = sum(self._market.companies[company_id].own_stock_ownership.values())
+        available_shares = self.total_shares - total_shares_owned
 
-        if self.cash < total_value:
-            return "Insufficient funds to buy stocks."
+        # Determine the maximum shares that can be bought based on available cash and available shares
+        max_affordable_shares = int(self.cash / self._market.get_stock_price(company_id))
+        amount = min(amount, available_shares, max_affordable_shares)
 
-        # Update stock ownership
-        if company_id == self.player_id:
-            # Buying own company's stock
-            self.own_stock_ownership[self.player_id] += amount
-        else:
-            # Buying another company's stock
-            if is_ai:
-                # For AI
-                self.stock_holdings[company_id] += amount
-            else:
-                # For player
-                self.stock_holdings[company_id] += amount
+        if amount == 0:
+            return "Cannot buy any shares due to insufficient funds or no available shares."
+
+        total_value = self._market.get_stock_price(company_id) * amount
+
+        self.cash -= total_value
+        self._market.update_stock_ledger(company_id, self.player_id, amount)
+
+        # Update buyer's stock_holdings
+        self.stock_holdings[company_id] += amount
+
+        # Update the target company's own_stock_ownership to reflect the buyer's new ownership
+        target_company = self._market.companies[company_id]
+        target_company.own_stock_ownership[self.player_id] = target_company.own_stock_ownership.get(self.player_id,
+                                                                                                    0) + amount
 
-        self.cash -= total_value  # Deduct the cost from the buyer's cash
         return f"Bought {amount} stocks of {company_id}."
 
     def _sell_stock(self, company_id, amount, total_value, is_ai):
-        """Handles the selling of stocks for the specified company."""
-        # Check if the seller has enough stocks to sell
-        stock_ownership = self._get_stock_ownership(company_id) if not is_ai else self.own_stock_ownership[company_id]
-        if stock_ownership < amount:
-            return "Not enough stocks to sell."
-
-        # Update stock ownership
-        if is_ai:
-            if company_id == self.player_id:
-                # AI selling its own stock
-                self.own_stock_ownership[company_id] -= amount
-            else:
-                # AI selling other company's stock
-                self.stock_holdings[company_id] -= amount
+        if not is_ai:
+            if self.stock_holdings[company_id] < amount:
+                return "Not enough stocks to sell."
+            # Update buyer's stock_holdings for non-AI (like player)
+            self.stock_holdings[company_id] -= amount
         else:
-            # Player selling stock
+            if self.own_stock_ownership[company_id] < amount:
+                return "Not enough stocks to sell."
+            # Update own_stock_ownership for AI
             self.own_stock_ownership[company_id] -= amount
 
-        self.cash += total_value  # Add the proceeds to the seller's cash
+        self.cash += total_value
+        self._market.update_stock_ledger(company_id, self.player_id, -amount)
+
+        # If an AI company is selling its own shares, update the ownership for all shareholders
+        if is_ai and company_id == self.player_id:
+            for shareholder in self._market.companies[company_id].own_stock_ownership.keys():
+                self._market.companies[company_id].own_stock_ownership[shareholder] -= amount * \
+                                                                                       (self._market.companies[
+                                                                                            company_id].own_stock_ownership[
+                                                                                            shareholder] / self.total_shares)
+
         return f"Sold {amount} stocks of {company_id}."
 
     def _calculate_available_shares(self):