nanocode/fumamxCRM/backend/fumamx_cli.py

"""
fumamxCLI 调用封装
"""
import subprocess
import json
import os
import re
from typing import List, Dict, Optional

CLI_PATH = "/home/ubuntu/nanocode/fumamxCLI/fumamx"
COOKIE_FILE = "/home/ubuntu/fumamx_cookies.txt"


def run_cli_command(args: List[str]) -> Dict:
    """运行 fumamxCLI 命令"""
    cmd = [CLI_PATH] + args
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=30,
            cwd="/home/ubuntu/nanocode/fumamxCLI"
        )
        
        # 提取 JSON 输出
        output = result.stdout
        
        # 尝试解析 JSON
        try:
            # 查找 JSON 部分（通常在 ╭─ 和 ╰─ 之间）
            json_match = re.search(r'[╭┌].*?[┐╮]([\s\S]*?)[╰└].*?[┘╯]', output)
            if json_match:
                json_str = json_match.group(1).strip()
                # 清理边框字符 (│ ║ ┃ ┆ ┇ etc.) 和 ANSI 转义序列
                json_str = re.sub(r'[│║┃┆┇┊┋]', '', json_str)
                json_str = re.sub(r'\x1b\[[0-9;]*m', '', json_str)
                json_str = json_str.strip()
                return json.loads(json_str)
            
            # 直接尝试解析整个输出（用于 --raw 模式）
            # 先清理 ANSI 转义序列
            cleaned_output = re.sub(r'\x1b\[[0-9;]*m', '', output)
            # 查找 JSON 对象（从第一个 { 到最后一个 }）
            json_start = cleaned_output.find('{')
            json_end = cleaned_output.rfind('}')
            if json_start != -1 and json_end != -1 and json_end > json_start:
                json_str = cleaned_output[json_start:json_end+1]
                return json.loads(json_str)
            
            return json.loads(cleaned_output)
        except json.JSONDecodeError:
            return {
                "success": result.returncode == 0,
                "output": output,
                "error": result.stderr
            }
            
    except subprocess.TimeoutExpired:
        return {"error": "命令执行超时"}
    except Exception as e:
        return {"error": str(e)}


def search_companies(keyword: str, page: int = 1, page_size: int = 10) -> Dict:
    """搜索公司
    
    返回: {
        "total": int,
        "results": [{...}],
        "keyword": str
    }
    """
    result = run_cli_command(["company", "search", keyword, "--page", str(page), "--page-size", str(page_size), "--raw"])
    
    if result.get("code") == "0" and "data" in result:
        data = result["data"]
        return {
            "total": int(data.get("total", 0)),
            "results": data.get("results", []),
            "keyword": keyword
        }
    
    return {"total": 0, "results": [], "keyword": keyword}


def get_company_detail(name: str, country: str = "") -> Dict:
    """获取公司详情（进口+出口数据）"""
    # 获取进口数据
    import_result = run_cli_command(["import-", "enterprise", name, "--country", country])
    
    # 获取出口数据
    export_result = run_cli_command(["export", "enterprise", name, "--country", country])
    
    # 获取统计
    import_stats = run_cli_command(["import-", "stats", name])
    export_stats = run_cli_command(["export", "stats", name])
    
    return {
        "name": name,
        "country": country,
        "import_data": import_result if import_result.get("code") == "0" else None,
        "export_data": export_result if export_result.get("code") == "0" else None,
        "import_stats": import_stats if import_stats.get("code") == "0" else None,
        "export_stats": export_stats if export_stats.get("code") == "0" else None,
    }


def update_cookie(cookie_content: str) -> bool:
    """更新 cookie 文件"""
    try:
        with open(COOKIE_FILE, 'w') as f:
            f.write(cookie_content)
        return True
    except Exception as e:
        print(f"保存 cookie 失败: {e}")
        return False


def check_cookie_valid() -> bool:
    """检查 cookie 是否有效"""
    result = run_cli_command(["status"])
    return result.get("code") == "0" or "API 连接正常" in str(result)


def get_countries() -> List[Dict]:
    """获取国家列表"""
    result = run_cli_command(["country", "list"])
    
    if result.get("code") == "0" and "data" in result:
        return result["data"]
    
    return []


class FumamxCLIWrapper:
    """fumamxCLI 包装类 - 直接方法调用"""
    
    def search_company(self, keyword: str, page: int = 1, page_size: int = 10) -> Dict:
        """搜索公司"""
        return run_cli_command([
            "company", "search", keyword,
            "--page", str(page),
            "--page-size", str(page_size),
            "--raw"
        ])
    
    def get_company_info(self, company_id: str, domain: str) -> Dict:
        """获取公司详细信息"""
        return run_cli_command([
            "company", "info", company_id, domain,
            "--raw"
        ])
    
    def get_company_contacts(self, company_id: str, domain: str) -> Dict:
        """获取公司联系人列表"""
        return run_cli_command([
            "company", "contacts", company_id, domain,
            "--raw"
        ])
    
    def get_company_phones(self, domain: str) -> Dict:
        """获取公司电话/图片信息"""
        return run_cli_command([
            "company", "phones", domain,
            "--raw"
        ])