from dotenv import load_dotenv
from crewai import Agent
from crewai.tools import BaseTool

from crew.tasks.data_access_tasks import create_data_collection_task

from crew.tools.internal_db_tools import (
    get_customer_details, get_credit_score,
    get_account_status, get_customer_id_by_name,
)
from crew.tools.government_db_tools import get_pr_status

load_dotenv()

"""
Tool wrapper
These classes wrap logic in try-except so that if tool fails,
it return a string error message instead of crashing the entire agent loop
"""

class GetCustomerIdByNameTool(BaseTool):
    name: str = "get_customer_id_by_name"
    description: str = (
        "Search for a customer's unique ID using their full name. "
        "Useful for converting a name into an ID. "
        "Returns the ID string if found, or 'None' if not found."
    )
    def _run(self, name: str):
        try:
            # strip away whitespaces
            clean_name = name.strip() 
            result = get_customer_id_by_name(name=clean_name)
            return result if result else "None"
        except Exception as e:
            return f"Error querying ID: {str(e)}"

class GetCustomerDetailsTool(BaseTool):
    name: str = "get_customer_details"
    description: str = "Fetch a customer's profile (email, nationality) using their customer_id."
    def _run(self, customer_id: str):
        try:
            return get_customer_details(customer_id=customer_id)
        except Exception as e:
            return f"Error fetching details: {str(e)}"

class GetCreditScoreTool(BaseTool):
    name: str = "get_credit_score"
    description: str = "Retrieve the numeric credit score (300-850) using a customer_id."
    def _run(self, customer_id: str):
        try:
            return get_credit_score(customer_id=customer_id)
        except Exception as e:
            return f"Error fetching score: {str(e)}"

class GetAccountStatusTool(BaseTool):
    name: str = "get_account_status"
    description: str = "Check if the account is Active, Closed, or Delinquent using a customer_id."
    def _run(self, customer_id: str):
        try:
            return get_account_status(customer_id=customer_id)
        except Exception as e:
            return f"Error fetching status: {str(e)}"

class GetPRStatusTool(BaseTool):
    name: str = "get_pr_status"
    description: str = "Check government PR/Residency status using a customer_id."
    def _run(self, customer_id: str):
        try:
            return get_pr_status(customer_id=customer_id)
        except Exception as e:
            return f"Error fetching PR status: {str(e)}"


"""
Database Agent
Simulate a 'microservice' environment where data have to be fetch from different API endpoint
Identity service (Look up customer ID using name as user will hardly provide customer id in their query)
Profile service (Look for customer detail)
Credit service (Credit score)
Gov Service (PR status)

Force the database agent to return in JSON format in a form of report

There are a chain of thought portion to ensure that database agent only query from the internal and government database
"""

class DataAccessAgent:
    def __init__(self, llm):
        self.agent = Agent(
            role="Senior Data Investigator",
            goal="Orchestrate multiple API calls to build a complete customer profile.",
            
            backstory=(
                "You are the **Senior Data Investigator**.\n"
                "You operate in a **Microservices Environment**. Data is fragmented across different systems.\n"
                "**YOUR JOB**: You are the 'Aggregator'. You must chain API calls together to build a full picture.\n\n"
                "**Policy search is NOT YOUR JOB**: Do not try to access database if it is a policy related question\n\n"
                
                "### YOUR API PROTOCOL (CHAIN OF THOUGHT):\n"
                "1. **Step 1 (The Key)**: Always start by calling `get_customer_id_by_name`. You cannot do anything without the ID.\n"
                "   - If ID is 'None', STOP immediately. Do not call other tools.\n"
                "2. **Step 2 (The loop)**: Once you have the `customer_id`, you must call the other 4 endpoints individually:\n"
                "   - `get_customer_details` (Email, Nationality)\n"
                "   - `get_credit_score` (Financial Health)\n"
                "   - `get_account_status` (Bank Standing)\n"
                "   - `get_pr_status` (Government Verification)\n"
                "3. **Step 3 (The Report)**: Consolidate all 5 results into one JSON.\n"

                "### 🛑 TOOL USAGE DECREE:\n"
                "1. **INPUT IS PLAIN TEXT**: Your tool input must be a simple, continuous string. Do not use { } or [ ].\n"
                "2. **NO NESTED KEYS**: Never use the word 'description' or 'query' inside your action input.\n"
                "3. **TRANSLATE TO PROSE**: If you received JSON data from a coworker, describe that data in a sentence when passing it to the next person.\n"
                "4. **CLEAN STRINGS**: Do not use backslashes (\), quotes inside quotes, or markdown (```) in tool calls."

                
            ),
            llm=llm,
            verbose=True,
            max_iter=10, 
            tools=[
                GetCustomerIdByNameTool(), 
                GetCustomerDetailsTool(), 
                GetCreditScoreTool(), 
                GetAccountStatusTool(), 
                GetPRStatusTool()
            ],
            allow_delegation=False
        )

    def get_task(self, query):
        return create_data_collection_task(self.agent, query)