# -*- coding: utf-8 -*-
import unittest
from unittest import mock

from flask import Flask, session

from routes import ai
from services.ai_datasources import AI_DATASOURCES


class AiAnalysisP0Test(unittest.TestCase):
    def setUp(self):
        self.app = Flask(__name__)
        self.app.secret_key = 'test'

    def test_normalize_tasks_accepts_only_registered_controlled_tools(self):
        raw = [
            {'datasource': 'product_profile', 'params': {'product_code': 'SKU-1', 'sql': 'DROP TABLE x', 'limit': 999}},
            {'datasource': 'product_profile', 'params': {'brand': '重复'}},
            {'datasource': 'not_exists', 'params': {}},
            {'datasource': 'planning_metrics', 'params': {'limit': 'bad'}},
        ]
        tasks = ai._normalize_tasks(raw)
        self.assertEqual(['product_profile', 'planning_metrics'], [x['datasource'] for x in tasks])
        self.assertNotIn('sql', tasks[0]['params'])
        self.assertEqual(200, tasks[0]['params']['limit'])
        self.assertEqual(50, tasks[1]['params']['limit'])

    def test_chart_must_reference_real_dataset_fields(self):
        datasets = [{'key': 'planning_metrics', 'columns': ['product_code', '建议下单数量']}]
        valid = ai._validate_chart({
            'type': 'bar', 'datasource': 'planning_metrics', 'label_field': 'product_code',
            'value_fields': ['建议下单数量'], 'title': '建议下单排行'
        }, datasets)
        self.assertEqual('bar', valid['type'])
        invalid = ai._validate_chart({
            'type': 'bar', 'datasource': 'planning_metrics', 'label_field': '不存在',
            'value_fields': ['建议下单数量']
        }, datasets)
        self.assertEqual('none', invalid['type'])

    def test_multi_source_query_then_second_stage_analysis(self):
        plan = {
            'intent': 'analysis', 'title': 'SKU下单风险分析',
            'tasks': [
                {'datasource': 'product_profile', 'params': {'product_code': 'SKU-1'}, 'purpose': '基础资料'},
                {'datasource': 'planning_metrics', 'params': {'product_code': 'SKU-1'}, 'purpose': '下单判断'},
            ]
        }
        analysis = {
            'title': 'SKU下单风险分析', 'summary': '当前建议暂缓下单。',
            'findings': [{'title': '库存', 'detail': '总支撑天数为40天', 'level': '信息'}],
            'risks': [], 'actions': [{'priority': 'P0', 'action': '核对在途', 'reason': '避免重复下单'}],
            'evidence': [{'label': '总支撑天数', 'value': '40天', 'source': '计划预测'}],
            'confidence': '高', 'limitations': [], 'follow_ups': ['查看BOM包材是否够用'],
            'chart': {'type': 'bar', 'datasource': 'planning_metrics', 'label_field': 'product_code',
                      'value_fields': ['建议下单数量'], 'title': '建议下单'}
        }
        fake_results = {
            'product_profile': {'columns': ['product_code', '供应商'], 'rows': [{'product_code': 'SKU-1', '供应商': 'A'}]},
            'planning_metrics': {'columns': ['product_code', '建议下单数量'], 'rows': [{'product_code': 'SKU-1', '建议下单数量': 0}]},
        }
        old_handlers = {key: AI_DATASOURCES[key]['handler'] for key in fake_results}
        try:
            for key, result in fake_results.items():
                AI_DATASOURCES[key]['handler'] = mock.Mock(return_value=result)
            with self.app.test_request_context('/api/ai/ask', method='POST', json={'question': 'SKU-1要不要下单'}):
                session['user_id'] = 1
                session['role'] = 'super_admin'
                with mock.patch.object(ai, 'operation_required', return_value=True), \
                     mock.patch.object(ai, '_plan_analysis', return_value=plan), \
                     mock.patch.object(ai, '_synthesize_analysis', return_value=analysis), \
                     mock.patch.object(ai, '_can_view_datasource', return_value=True), \
                     mock.patch.object(ai, 'datasource_meta', side_effect=lambda key: {'key': key, 'name': AI_DATASOURCES[key]['name'], 'data_date': '2026-10-01'}), \
                     mock.patch.object(ai, '_save_history') as save_history, \
                     mock.patch.object(ai, 'add_operation_log'):
                    response = ai.ai_ask.__wrapped__()
            payload = response.get_json()['data']
            self.assertEqual('analysis', payload['type'])
            self.assertEqual(2, len(payload['datasets']))
            self.assertEqual('当前建议暂缓下单。', payload['summary'])
            self.assertEqual('bar', payload['chart']['type'])
            self.assertTrue(save_history.called)
            self.assertEqual('analysis', save_history.call_args.args[4])
        finally:
            for key, handler in old_handlers.items():
                AI_DATASOURCES[key]['handler'] = handler

    def test_p0_datasources_cover_product_bom_order_and_supplier(self):
        for key in ('product_profile', 'bom_supply_chain', 'order_detail', 'supplier_delivery'):
            self.assertIn(key, AI_DATASOURCES)
            self.assertTrue(callable(AI_DATASOURCES[key]['handler']))
            self.assertTrue(AI_DATASOURCES[key].get('permission_pages'))


if __name__ == '__main__':
    unittest.main()