from typing import Optional
import pandas as pd
from langchain_openai import ChatOpenAI
import warnings

from src.PostgresSQL.db_handler import PgHandler
from src.config.load_config import project_config

LLM = ChatOpenAI(
    model="qwen-plus-latest",
    api_key=project_config.get('bailian_API_key', ''),
    base_url=project_config.get('bailian_API_URL', ''),
)


def load_history_by_year_month(year: int, month: Optional[int] = None,
                               table_name: str = '王适意_personal_chat',
                               schema: str = 'public',
                               db_config: Optional[dict] = None,
                               week: Optional[int] = None) -> pd.DataFrame:
    """从 Postgres 中读取历史消息并按年/月或年/周过滤。

    参数：
      - year: 年份（四位数）
      - month: 可选的月份（1-12），若提供则只返回该月的数据
      - week: 可选的每年的第几周（ISO week）。如果提供则优先按周过滤，会 bypass month，但会发出警告。
      - table_name: 数据库表名，默认 '王适意_personal_chat'
      - schema: schema，默认 'public'
      - db_config: 可选的数据库连接参数字典（含 host, port, user, password, database）

    返回：pandas.DataFrame，包含原始列并新增一列'dt'（datetime 类型）用于时间筛选与排序。
    """
    # 使用传入的 db_config 或 project_config
    cfg = db_config or project_config
    host = cfg.get('host') or cfg.get('db_host')
    port = int(cfg.get('port') or cfg.get('db_port') or 5432)
    user = cfg.get('user') or cfg.get('db_user')
    password = cfg.get('password') or cfg.get('db_password')
    database = cfg.get('tg_message_db_name')

    pg = PgHandler(host=host, port=port, user=user, password=password, database=database)
    try:
        if not pg.table_exists(table_name, schema=schema):
            raise ValueError(f"table {schema}.{table_name} not exists")
        df = pg.read_table(table_name, schema=schema)
    finally:
        pg.close()

    if df.empty:
        return df

    # 尝试解析时间列：优先 message_datetime，再尝试 message_unix
    dt = None
    if 'message_datetime' in df.columns:
        try:
            dt = pd.to_datetime(df['message_datetime'], errors='coerce')
        except Exception:
            dt = None
    if (dt is None or dt.isna().all()) and 'message_unix' in df.columns:
        # message_unix 可能为字符串或数值
        try:
            dt = pd.to_datetime(df['message_unix'].astype(float), unit='s', errors='coerce')
        except Exception:
            dt = pd.to_datetime(df['message_unix'], errors='coerce')

    if dt is None:
        # 无法解析时间，返回原表
        return df

    df = df.copy()
    df['dt'] = dt

    # 如果同时提供 month 和 week，警告并以 week 为准（week 优先）
    if week is not None and month is not None:
        warnings.warn('Both "month" and "week" provided; "week" will take precedence and "month" will be ignored.', UserWarning)

    # 过滤年份和可选月份/周
    # 当提供 week 时，按 ISO week 过滤，并且考虑 ISO week 的 year（isocalendar year）以确保正确匹配
    if week is not None:
        # Pandas >= 1.1 提供 dt.isocalendar()
        try:
            iso = df['dt'].dt.isocalendar()
            # iso is a DataFrame-like with 'year' and 'week' columns
            df = df[(iso['year'] == int(year)) & (iso['week'] == int(week))]
        except Exception:
            # 备选：使用 datetime.isocalendar() per-row（较慢）
            df = df[df['dt'].apply(lambda x: x.isocalendar()[0] == int(year) and x.isocalendar()[1] == int(week))]
    else:
        # 默认按 calendar 年过滤
        df = df[df['dt'].dt.year == int(year)]
        if month is not None:
            df = df[df['dt'].dt.month == int(month)]

    df = df.sort_values('dt').reset_index(drop=True)
    return df


def do_analysis_on_history_messages(df: pd.DataFrame):
    data = df.to_dict(orient='records')
    to_model_data = [
        {
            'msg_time': item.get('message_datetime'),
            'msg_from': item.get('message_from'),
            'msg_text': f"转发来自{item['forward_from']}的消息: {item.get('message_text')}" if item.get('forward_from') else item.get('message_text'),
            'file_type': item.get('file_type'),
        }
        for item in data
    ]
    to_model_data = to_model_data
    prompt = f"""
你是一个聊天记录分析专家，擅长根据用户的聊天记录，分析用户的行为，然后撰写报告。
你的报告应当包含以下内容：
1. 以天为单位，分析用户在这一天内的主要活动和重大事件。每一天的总结应当独立成段。每一天的重大事件应当按照时间顺序排列。
2. 输出总结报告。

你的总结应当包括但不限于以下内容：
- 用户在这个月内的主要活动和重大事件。
- 用户的情绪变化和心理状态。
- 用户的兴趣爱好和关注点。
- 用户的社交活动和人际关系。
- 用户的生活习惯和日常规律。
- 任何其他你认为重要的信息。

注意：不要对于进行过度的推断，只基于用户的聊天记录内容进行分析。

注意：用户是Z Armor，与他对话的是王适意，与用户是恋人关系。司，王适意，Skuld都是指王适意；装，Armor，周义人都是指Z Armor。一定要搞清楚双方的身份关系，message_from字段指的是，这个消息是谁发的。

以下是用户的日记内容：
{to_model_data}
    """
    llm_resp = LLM.invoke(prompt)
    print(llm_resp.content)
    return llm_resp.content


if __name__ == "__main__":
    # df = load_history_by_year_month(2025, month=1)
    df = load_history_by_year_month(2025, week=2)
    summary = do_analysis_on_history_messages(df)
    print(f"done")
