#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2023/5/12 00:30
@Author  : alexanderwu
@File    : software_company.py
@Modified By: mashenquan, 2023/11/27. Add an archiving operation after completing the project, as specified in
        Section 2.2.3.3 of RFC 135.
"""
from pathlib import Path
from pydantic import BaseModel, Field

from metagpt.actions import UserRequirement
from metagpt.config import CONFIG
from metagpt.const import MESSAGE_ROUTE_TO_ALL
from metagpt.environment import Environment
from metagpt.logs import logger
from metagpt.roles import Role
from metagpt.schema import Message
from metagpt.utils.common import NoMoneyException
from metagpt.utils.utils import read_json_file, write_json_file


class Team(BaseModel):
    """
    Team: Possesses one or more roles (agents), SOP (Standard Operating Procedures), and a platform for instant messaging,
    dedicated to perform any multi-agent activity, such as collaboratively writing executable code.
    """

    env: Environment = Field(default_factory=Environment)
    investment: float = Field(default=10.0)
    idea: str = Field(default="")

    class Config:
        arbitrary_types_allowed = True

    def serialize(self, stg_path: Path):
        team_info_path = stg_path.joinpath("team_info.json")
        write_json_file(team_info_path, {
            "idea": self.idea,
            "investment": self.investment
        })

        self.environment.serialize(stg_path.joinpath("environment"))

    def deserialize(self, stg_path: Path):
        """ stg_path = ./storage/team """
        # recover team_info
        team_info_path = stg_path.joinpath("team_info.json")
        if not team_info_path.exists():
            logger.error("recover storage not exist, not to recover and continue run the old project.")
        team_info = read_json_file(team_info_path)
        self.investment = team_info.get("investment", 10.0)
        self.idea = team_info.get("idea", "")

        # recover environment
        environment_path = stg_path.joinpath("environment")
        self.environment = Environment()
        self.environment.deserialize(stg_path=environment_path)

    def hire(self, roles: list[Role]):
        """Hire roles to cooperate"""
        self.env.add_roles(roles)

    def invest(self, investment: float):
        """Invest company. raise NoMoneyException when exceed max_budget."""
        self.investment = investment
        CONFIG.max_budget = investment
        logger.info(f"Investment: ${investment}.")

    def _check_balance(self):
        if CONFIG.total_cost > CONFIG.max_budget:
            raise NoMoneyException(CONFIG.total_cost, f"Insufficient funds: {CONFIG.max_budget}")

    def run_project(self, idea, send_to: str = ""):
        """Start a project from publishing user requirement."""
        self.idea = idea

        # Human requirement.
        self.env.publish_message(
            Message(role="Human", content=idea, cause_by=UserRequirement, send_to=send_to or MESSAGE_ROUTE_TO_ALL)
        )

    def _save(self):
        logger.info(self.json(ensure_ascii=False))

    async def run(self, n_round=3):
        """Run company until target round or no money"""
        while n_round > 0:
            # self._save()
            n_round -= 1
            logger.debug(f"max {n_round=} left.")
            self._check_balance()
            await self.env.run()
        if CONFIG.git_repo:
            CONFIG.git_repo.archive()
        return self.env.history