import unittest
import os

from omegaconf import OmegaConf
from vectara_agent.agent import Agent

import sqlite3
from datasets import load_dataset

from app import setup_db
from agent import initialize_agent

from dotenv import load_dotenv
load_dotenv(override=True)

class TestAgentResponses(unittest.TestCase):

    def test_responses(self):

        cfg = OmegaConf.create({
            'customer_id': str(os.environ['VECTARA_CUSTOMER_ID']),
        'corpus_ids': str(os.environ['VECTARA_CORPUS_IDS']).split(','),
        'api_keys': str(os.environ['VECTARA_API_KEYS']).split(','),
        'examples': os.environ.get('QUERY_EXAMPLES', None)
        })

        setup_db()

        agent = initialize_agent(_cfg=cfg)
        self.assertIsInstance(agent, Agent)

        # Knows types of electric vehicles
        type_output = agent.chat('What are the different types of electric vehicles? Only provide the name of each type, nothing else.').lower()
        ev_types = ['battery', 'hybrid', 'plug-in hybrid', 'fuel cell']

        for ev_type in ev_types:
            self.assertIn(ev_type, type_output)


        # Questions about car models - ev query tool
        self.assertIn('mach-e', agent.chat('Which EV is made by Mustang? Provide the name of the model only.').lower())
        self.assertIn('fuel cell', agent.chat('What EV type is the Toyota Mirai? Just give the type name.').lower())
        
        
        # Incentive query tool tests
        self.assertIn('no', agent.chat('Does the U.S. Department of Defense offer incentives for purchasing electric vehicles? Only say "yes" or "no".').lower())
        self.assertIn('2035', agent.chat('At what year must all new passenger vehicles be zero emission vehicles in California? Give the year only.').lower())


        # Database Tool questions
        self.assertIn('king', agent.chat('Which county in the state of Washington had the highest number of EV registrations in 2023? Provide the name only.').lower())
        self.assertIn('seattle', agent.chat('Which city in the state of Washington had the highest number of EV registrations in 2023? Provide the name only.').lower())
        self.assertIn('tesla model y', agent.chat('What car was the most popular in Seattle in 2023? Provide the make and model only').lower())

        # Misc. questions
        self.assertIn('tesla', agent.chat('Which company developed a standard charging port for electric cars in California? Only provide the company name, nothing else').lower())

if __name__ == "__main__":
    setup_db()
    unittest.main()