import unittest

from app.services.company_storage import extract_company_fields, normalize_company_id


class CompanyStorageTests(unittest.TestCase):
    def test_normalize_qcwy_company_id(self):
        self.assertEqual(normalize_company_id("qcwy", "co123"), "123")
        self.assertEqual(normalize_company_id("qcwy", "123"), "123")
        self.assertEqual(normalize_company_id("boss", "co123"), "co123")

    def test_extract_boss_fields(self):
        payload = {
            "zpData": {
                "brandComInfoVO": {
                    "encryptBrandId": "boss-1",
                    "brandName": "Boss公司",
                    "industryName": "互联网",
                    "scaleName": "100-499人",
                    "stageName": "B轮",
                    "logo": "https://example.com/logo.png",
                    "introduce": "品牌简介",
                },
                "companyFullInfoVO": {
                    "name": "Boss公司",
                    "typeName": "民营",
                    "cityName": "上海",
                    "address": "上海市徐汇区",
                    "website": "https://boss.example.com",
                },
            }
        }
        result = extract_company_fields("boss", payload, "boss-1")
        self.assertEqual(result["source_company_id"], "boss-1")
        self.assertEqual(result["company_name"], "Boss公司")
        self.assertEqual(result["industry"], "互联网")
        self.assertEqual(result["financing_stage"], "B轮")

    def test_extract_qcwy_fields(self):
        payload = {
            "coinfo": {
                "coid": "123",
                "coname": "前程公司",
                "cotype": "民营",
                "indtype1": "制造业",
                "cosize": "500-999人",
                "areaString": "广州",
                "caddr": "广州市天河区",
                "webUrl": "https://qcwy.example.com",
                "logourl": "https://qcwy.example.com/logo.png",
                "coinfo": "公司简介",
            },
            "financingStage": {"name": "未融资"},
        }
        result = extract_company_fields("qcwy", payload, "co123")
        self.assertEqual(result["source_company_id"], "123")
        self.assertEqual(result["company_name"], "前程公司")
        self.assertEqual(result["company_size"], "500-999人")
        self.assertEqual(result["website"], "https://qcwy.example.com")

    def test_extract_zhilian_fields(self):
        payload = {
            "data": {
                "companyBase": {
                    "companyNumber": "zl-1",
                    "companyName": "智联公司",
                    "companyTypeName": "上市公司",
                    "industryName": "教育",
                    "companySize": "1000-9999人",
                    "financingStage": {"name": "已上市"},
                    "cityName": "北京",
                    "address": "北京市海淀区",
                    "companyUrl": "https://zl.example.com",
                    "logoUrl": "https://zl.example.com/logo.png",
                    "companyDescWithHtml": "<p>公司简介</p>",
                }
            }
        }
        result = extract_company_fields("zhilian", payload, "zl-1")
        self.assertEqual(result["source_company_id"], "zl-1")
        self.assertEqual(result["company_name"], "智联公司")
        self.assertEqual(result["company_type"], "上市公司")
        self.assertEqual(result["description"], "<p>公司简介</p>")


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