import requests
import os
from bs4 import BeautifulSoup
from PIL import Image
import io

class AccountSpider:
    def __init__(self, save_dir: str = None):
        self.base_url = "http://121.204.253.175:8088"
        self.session = requests.Session()
        self.save_dir = save_dir or os.path.dirname(os.path.abspath(__file__))
        # 设置默认请求头
        self.session.headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }

    def get_verification_code(self):
        """获取验证码图片，返回图片字节数据"""
        code_url = f"{self.base_url}/code.asp"
        response = self.session.get(code_url)
        
        # 保存验证码图片到本地
        image_path = os.path.join(self.save_dir, 'verification_code.png')
        image = Image.open(io.BytesIO(response.content))
        image.save(image_path)
        return response.content

    def login(self, username, password,v_code=""):
        """执行登录操作"""

        # 获取验证码
        if v_code:
            verification_code = v_code
        else:
            verification_code = self.get_verification_code()
        
        # 准备登录数据
        login_data = {
            'login_type': '0',
            'f_user': username,
            'f_code': password,
            'codeOK': verification_code,
            'Submit': '%C8%B7%B6%A8'
        }
        
        # 发送登录请求
        login_url = f"{self.base_url}/login_result.asp"
        response = self.session.post(login_url, data=login_data)
        response.encoding = 'gb2312'  # 设置正确的编码
        
        # 检查登录是否成功 - 通过检查是否包含重定向到account.asp的脚本
        if "window.location.href=\"account.asp\"" in response.text:
            return True
        return False

    def get_balance(self):
        """获取账户余额"""
        account_url = f"{self.base_url}/account.asp"
        response = self.session.get(account_url)
        response.encoding = 'gb2312'  # 设置正确的编码
        
        soup = BeautifulSoup(response.text, 'html.parser')
        balance_text = soup.find_all('span', class_='red')[1].text
        return float(balance_text)

def main():
    """仅用于独立测试，实际使用通过 nonebot 插件调用"""
    import os
    username = os.environ.get("DAMO_USERNAME", "")
    password = os.environ.get("DAMO_PASSWORD", "")
    if not username or not password:
        print("请设置环境变量 DAMO_USERNAME 和 DAMO_PASSWORD")
        return
    
    spider = AccountSpider()
    
    if spider.login(username, password):
        print("登录成功！")
        balance = spider.get_balance()
        print(f"账户余额：{balance}元")
    else:
        print("登录失败，请检查账号密码或验证码是否正确")