import json
import os
from typing import Dict, Any


class OutputWriter:
    """将 AI 生成的数据写入文件系统。"""

    def __init__(self, output_dir: str):
        self.output_dir = output_dir

    def write(self, program_id: str, ai_result: Dict[str, Any],
              input_type: str) -> Dict[str, str]:
        """写入所有组的输出文件。

        Returns:
            {group_folder: written_file_path} 映射
        """
        written = {}
        groups = ai_result.get('groups', ai_result)

        for group_key in sorted(groups.keys()):
            group_data = groups[group_key]
            group_dir = os.path.join(self.output_dir, program_id, group_key)
            os.makedirs(group_dir, exist_ok=True)

            if input_type in ('file', 'mixed'):
                json_path = self._write_json(group_dir, program_id, group_key, group_data)
                written[f"{group_key}/json"] = json_path

            if input_type in ('db', 'mixed'):
                sql_path = self._write_sql(group_dir, program_id, group_key, group_data)
                written[f"{group_key}/sql"] = sql_path

        return written

    def _write_json(self, group_dir: str, program_id: str,
                    group_key: str, data: Any) -> str:
        """写入 JSON 文件。"""
        filename = f"{program_id}_{group_key}.json"
        filepath = os.path.join(group_dir, filename)

        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

        return filepath

    def _write_sql(self, group_dir: str, program_id: str,
                   group_key: str, data: Any) -> str:
        """写入 SQL 文件。"""
        filename = f"{program_id}_{group_key}.sql"
        filepath = os.path.join(group_dir, filename)

        sql_content = data if isinstance(data, str) else data.get('sql', json.dumps(data, ensure_ascii=False))

        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(sql_content)

        return filepath