from __future__ import annotations

import importlib.util
import json
import tempfile
import unittest
from pathlib import Path

from src.agent_tools import build_tool_context, default_tool_registry, execute_tool
from src.agent_types import AgentRuntimeConfig
from src.data_agent_router_sessions import (
    convert_router_candidates_to_records,
    profile_router_sessions,
    sample_router_candidates,
    search_router_sessions,
)


HAS_PYARROW = importlib.util.find_spec('pyarrow') is not None


class DataAgentRouterSessionTests(unittest.TestCase):
    def test_sample_router_candidates_accepts_search_payload(self) -> None:
        payload = {
            'candidates': [
                _candidate('s1', 'phone', 'QA', 'unknown', '怎么查找设备'),
                _candidate('s2', 'car', 'mapCopilot', '导航', '导航去公司'),
                _candidate('s3', 'phone', 'QA', 'unknown', '手机丢了怎么办'),
            ]
        }

        result = sample_router_candidates(payload, sample_size=2, strategy='stride')

        self.assertEqual(result['total_candidate_count'], 3)
        self.assertEqual(result['sampled_count'], 2)
        self.assertEqual(result['summary']['device'][0], {'value': 'phone', 'count': 2})
        self.assertEqual(result['candidates'][0]['semantic_session_id'], 's1')
        self.assertEqual(result['candidates'][1]['semantic_session_id'], 's3')

    def test_sample_router_candidates_accepts_json_string(self) -> None:
        result = sample_router_candidates(
            json.dumps([_candidate('s1', 'phone', 'QA', 'unknown', '怎么查找设备')], ensure_ascii=False),
            sample_size=1,
        )

        self.assertEqual(result['sampled_count'], 1)
        self.assertEqual(result['candidates'][0]['matched_turn']['query'], '怎么查找设备')

    def test_convert_router_candidates_to_records_keeps_online_metadata(self) -> None:
        candidate = _candidate('s1', 'phone', 'QA', 'unknown', '总结一下')
        candidate['turn_count'] = 2
        candidate['matched_turn_index'] = 1
        candidate['matched_turn']['timestamp'] = 20
        candidate['prev_turns'] = [{'timestamp': 10, 'query': '这篇文章讲了什么', 'tts': ''}]

        result = convert_router_candidates_to_records(
            [candidate],
            dataset_label='总结类边界评测集',
            default_target='Summarize',
            default_complex=True,
            batch_id='summary',
        )

        self.assertEqual(result['record_count'], 1)
        record = result['records'][0]
        self.assertEqual(record['source']['type'], 'online')
        self.assertEqual(record['source']['request_id'], 's1-rid')
        self.assertEqual(record['turn']['query'], '总结一下')
        self.assertEqual(record['prev_session'], [{'query': '这篇文章讲了什么', 'tts': '', 'timestamp': 10}])
        self.assertEqual(record['context']['domain'], 'QA')
        self.assertEqual(record['label']['target'], 'Summarize')
        self.assertEqual(record['dimensions'], {'complex': True})

    def test_convert_router_candidates_to_records_applies_review_decisions(self) -> None:
        candidates = [
            _candidate('s1', 'phone', 'QA', 'unknown', '总结一下'),
            _candidate('s2', 'phone', 'QA', 'unknown', '总结一下红楼梦'),
        ]

        result = convert_router_candidates_to_records(
            candidates,
            dataset_label='总结类边界评测集',
            review_decisions=[
                {
                    'semantic_session_id': 's1',
                    'matched_turn_index': 0,
                    'decision': 'include',
                    'target': 'Summarize',
                    'complex': True,
                    'notes': '前文有可总结内容',
                },
                {
                    'semantic_session_id': 's2',
                    'matched_turn_index': 0,
                    'decision': 'exclude',
                    'notes': '开放知识对象',
                },
            ],
        )

        self.assertEqual(result['record_count'], 1)
        self.assertEqual(result['skipped_count'], 1)
        self.assertEqual(result['records'][0]['meta']['notes'], '前文有可总结内容')
        self.assertEqual(result['records'][0]['dimensions']['complex'], True)

    @unittest.skipUnless(HAS_PYARROW, 'pyarrow is required for parquet tests')
    def test_profile_and_search_router_sessions_read_parquet(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            _write_router_parquet(Path(tmp_dir))

            profile = profile_router_sessions(tmp_dir, dates=['20260428'], max_files=1)
            search = search_router_sessions(
                tmp_dir,
                dates=['20260428'],
                domains=['mapCopilot'],
                query_keywords=['导航'],
                max_files=1,
                max_candidates=5,
            )

        self.assertEqual(profile['sampled_row_count'], 2)
        self.assertEqual(profile['missing_req_id_count'], 0)
        self.assertTrue(any(item['value'] == 'phone' for item in profile['distributions']['device']))
        self.assertEqual(search['candidate_count'], 1)
        candidate = search['candidates'][0]
        self.assertEqual(candidate['req_id'], 'rid-1')
        self.assertEqual(candidate['matched_turn']['query'], '导航去公司')
        self.assertEqual(candidate['prev_turns'][0]['query'], '你好小爱')

    @unittest.skipUnless(HAS_PYARROW, 'pyarrow is required for parquet tests')
    def test_router_session_tools_allow_external_absolute_paths(self) -> None:
        with tempfile.TemporaryDirectory() as workspace_dir, tempfile.TemporaryDirectory() as data_dir:
            _write_router_parquet(Path(data_dir))
            external_partition = Path(data_dir) / 'router_session_parquet' / 'date=20260428'

            profile = profile_router_sessions(
                workspace_dir,
                paths=[str(external_partition)],
                max_files=1,
            )
            search = search_router_sessions(
                workspace_dir,
                paths=[str(external_partition)],
                query_keywords=['导航'],
                max_files=1,
                max_candidates=5,
            )

        self.assertEqual(profile['sampled_row_count'], 2)
        self.assertEqual(search['candidate_count'], 1)
        self.assertEqual(search['candidates'][0]['matched_turn']['query'], '导航去公司')

    @unittest.skipUnless(HAS_PYARROW, 'pyarrow is required for parquet tests')
    def test_router_session_tools_execute_against_registry(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            _write_router_parquet(Path(tmp_dir))
            context = build_tool_context(AgentRuntimeConfig(cwd=Path(tmp_dir)))
            registry = default_tool_registry()

            search_result = execute_tool(
                registry,
                'data_agent_search_router_sessions',
                {'dates': ['20260428'], 'query_regex': '查找设备|手机丢了', 'max_files': 1},
                context,
            )
            self.assertTrue(search_result.ok, search_result.content)
            sample_result = execute_tool(
                registry,
                'data_agent_sample_router_candidates',
                {'candidates': json.loads(search_result.content), 'sample_size': 1, 'strategy': 'first'},
                context,
            )
            convert_result = execute_tool(
                registry,
                'data_agent_convert_router_candidates_to_records',
                {
                    'candidates': json.loads(sample_result.content),
                    'dataset_label': '查找设备评测集',
                    'default_target': 'QA',
                    'batch_id': 'find_device',
                },
                context,
            )

        self.assertTrue(sample_result.ok, sample_result.content)
        self.assertEqual(json.loads(sample_result.content)['sampled_count'], 1)
        self.assertTrue(convert_result.ok, convert_result.content)
        self.assertTrue(json.loads(convert_result.content)['validation']['ok'])


def _candidate(session_id: str, device: str, domain: str, intent: str, query: str) -> dict[str, object]:
    return {
        'semantic_session_id': session_id,
        'req_id': f'{session_id}-rid',
        'device': device,
        'turn_count': 1,
        'matched_turn_index': 0,
        'matched_turn': {
            'timestamp': 1,
            'query': query,
            'tts': '',
            'action': {'domain': domain, 'intent': intent},
        },
        'prev_turns': [],
        'source_path': 'router_session_parquet/date=20260428/part-0.parquet',
        'date': '20260428',
    }


def _write_router_parquet(root: Path) -> None:
    import pyarrow as pa
    import pyarrow.parquet as pq

    partition = root / 'router_session_parquet' / 'date=20260428'
    partition.mkdir(parents=True)
    rows = [
        {
            'semantic_session_id': 'phone_20260428_1',
            'device': 'phone',
            'req_id': 'rid-1',
            'turn_count': 2,
            'turns_array': [
                {
                    'timestamp': 1,
                    'query': '你好小爱',
                    'action_json': '{"domain":"dialogCopilot","intent":"问候"}',
                    'tts': '你好',
                },
                {
                    'timestamp': 2,
                    'query': '导航去公司',
                    'action_json': '{"domain":"mapCopilot","intent":"导航","func":"start_navigation"}',
                    'tts': '开始导航',
                },
            ],
        },
        {
            'semantic_session_id': 'phone_20260428_2',
            'device': 'phone',
            'req_id': 'rid-2',
            'turn_count': 1,
            'turns_array': [
                {
                    'timestamp': 3,
                    'query': '手机丢了怎么查找设备',
                    'action_json': '{"domain":"QA","intent":"unknown"}',
                    'tts': '',
                }
            ],
        },
    ]
    pq.write_table(pa.Table.from_pylist(rows), partition / 'part-0.parquet')


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