#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import json
import re
import sys
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Dict, Iterable, List, Optional


ASSESSMENT_LABELS = {
    "SCL90-2078": "SCL90",
    "SDS-2082": "SDS",
    "SAS-2083": "SAS",
}

RISK_ORDER = {
    "正常": 0,
    "低风险": 1,
    "中风险": 2,
    "高风险": 3,
}


@dataclass
class User:
    id: int
    email: str
    role: str
    enabled: bool
    created_at: str


@dataclass
class Profile:
    user_id: int
    name: str
    phone: str
    updated_at: str
    gender: str
    id_card: str
    birth_month: str
    education: str
    school_major: str


@dataclass
class Assessment:
    id: int
    test_paper_id: str
    title: str
    description: str
    active: bool


@dataclass
class Submission:
    id: int
    user_id: int
    assessment_id: int
    raw_score: int
    level: str
    dimension_scores: str
    submitted_at: str


@dataclass
class Report:
    id: int
    submission_id: int
    user_id: int
    summary: str
    suggestions: str
    created_at: str


@dataclass
class Question:
    id: int
    assessment_id: int
    code: str
    dimension: str
    prompt: str
    sort_order: int


@dataclass
class Answer:
    id: int
    submission_id: int
    question_id: int
    answer_value: int


@dataclass
class CompositeResult:
    completed_all: bool
    progress: str
    score: Optional[int]
    level: str


def parse_bool(value: str) -> bool:
    return value == "t"


def blank(value: Optional[str]) -> bool:
    return value is None or value.strip() == ""


def clamp(value: float, lower: float, upper: float) -> float:
    return max(lower, min(upper, value))


def max_level(level_a: str, level_b: str) -> str:
    return level_a if RISK_ORDER.get(level_a, -1) >= RISK_ORDER.get(level_b, -1) else level_b


def from_score(score: int) -> str:
    if score >= 75:
        return "高风险"
    if score >= 50:
        return "中风险"
    if score >= 25:
        return "低风险"
    return "正常"


def parse_json_int_field(raw: str, field: str) -> Optional[int]:
    if blank(raw):
        return None
    try:
        parsed = json.loads(raw)
    except Exception:
        return None
    value = parsed.get(field)
    if value is None:
        return None
    try:
        return int(value)
    except Exception:
        return None


def parse_scl_average_from_submission(submission: Optional[Submission]) -> float:
    if submission is None or blank(submission.dimension_scores):
        return 1.5
    scaled = parse_json_int_field(submission.dimension_scores, "总均分")
    if scaled is None:
        return 1.5
    return scaled / 100.0


def build_composite_result(scl: Optional[Submission], sds: Optional[Submission], sas: Optional[Submission]) -> CompositeResult:
    completed = (0 if scl is None else 1) + (0 if sds is None else 1) + (0 if sas is None else 1)
    progress = f"{completed}/3"
    if completed < 3:
        return CompositeResult(False, progress, None, "未完成")

    sds_score = sds.raw_score
    sas_score = sas.raw_score
    scl_avg = parse_scl_average_from_submission(scl)
    sds_risk = clamp((sds_score - 53) / 20.0 * 100.0, 0, 100)
    sas_risk = clamp((sas_score - 50) / 20.0 * 100.0, 0, 100)
    scl_risk = clamp((scl_avg - 1.5) / 3.0 * 100.0, 0, 100)
    score = round(0.30 * sds_risk + 0.30 * sas_risk + 0.40 * scl_risk)
    level = from_score(score)

    scl_positive = not blank(scl.level) and scl.level.strip().startswith("阳性")
    if sds_score >= 73 or sas_score >= 70:
        level = "高风险"
    if scl_positive and (sds_score >= 63 or sas_score >= 60):
        level = max_level(level, "中风险")
    if sds_score >= 63 and sas_score >= 60:
        level = max_level(level, "中风险")
    if scl_positive:
        level = max_level(level, "低风险")
    return CompositeResult(True, progress, score, level)


def iter_copy_blocks(sql_path: Path) -> Iterable[tuple[str, List[str], Iterable[List[str]]]]:
    header_re = re.compile(r"^COPY public\.([a-zA-Z0-9_]+) \(([^)]*)\) FROM stdin;$")
    with sql_path.open("r", encoding="utf-8", errors="replace") as handle:
        line = handle.readline()
        while line:
            match = header_re.match(line.rstrip("\n"))
            if not match:
                line = handle.readline()
                continue
            table = match.group(1)
            columns = [item.strip() for item in match.group(2).split(",")]

            def row_iter(source):
                for row in source:
                    row = row.rstrip("\n")
                    if row == r"\.":
                        break
                    yield row.split("\t")

            rows = list(row_iter(handle))
            yield table, columns, rows
            line = handle.readline()


def load_backup(sql_path: Path):
    users: Dict[int, User] = {}
    profiles: Dict[int, Profile] = {}
    assessments: Dict[int, Assessment] = {}
    submissions: Dict[int, Submission] = {}
    reports_by_submission: Dict[int, Report] = {}
    questions: Dict[int, Question] = {}
    answers_by_submission: Dict[int, List[Answer]] = defaultdict(list)

    for table, columns, rows in iter_copy_blocks(sql_path):
        index = {name: idx for idx, name in enumerate(columns)}
        if table == "users":
            for row in rows:
                user = User(
                    id=int(row[index["id"]]),
                    email=row[index["email"]],
                    role=row[index["role"]],
                    enabled=parse_bool(row[index["enabled"]]),
                    created_at=row[index["created_at"]],
                )
                users[user.id] = user
        elif table == "profiles":
            for row in rows:
                profile = Profile(
                    user_id=int(row[index["user_id"]]),
                    name=row[index["name"]],
                    phone=row[index["phone"]],
                    updated_at=row[index["updated_at"]],
                    gender=row[index["gender"]],
                    id_card=row[index["id_card"]],
                    birth_month=row[index["birth_month"]],
                    education=row[index["education"]],
                    school_major=row[index["school_major"]],
                )
                profiles[profile.user_id] = profile
        elif table == "assessments":
            for row in rows:
                assessment = Assessment(
                    id=int(row[index["id"]]),
                    test_paper_id=row[index["test_paper_id"]],
                    title=row[index["title"]],
                    description=row[index["description"]],
                    active=parse_bool(row[index["active"]]),
                )
                assessments[assessment.id] = assessment
        elif table == "submissions":
            for row in rows:
                submission = Submission(
                    id=int(row[index["id"]]),
                    user_id=int(row[index["user_id"]]),
                    assessment_id=int(row[index["assessment_id"]]),
                    raw_score=int(row[index["raw_score"]]),
                    level=row[index["level"]],
                    dimension_scores=row[index["dimension_scores"]],
                    submitted_at=row[index["submitted_at"]],
                )
                submissions[submission.id] = submission
        elif table == "reports":
            for row in rows:
                report = Report(
                    id=int(row[index["id"]]),
                    submission_id=int(row[index["submission_id"]]),
                    user_id=int(row[index["user_id"]]),
                    summary=row[index["summary"]],
                    suggestions=row[index["suggestions"]],
                    created_at=row[index["created_at"]],
                )
                reports_by_submission[report.submission_id] = report
        elif table == "questions":
            for row in rows:
                question = Question(
                    id=int(row[index["id"]]),
                    assessment_id=int(row[index["assessment_id"]]),
                    code=row[index["code"]],
                    dimension=row[index["dimension"]],
                    prompt=row[index["prompt"]],
                    sort_order=int(row[index["sort_order"]]),
                )
                questions[question.id] = question
        elif table == "answers":
            for row in rows:
                answer = Answer(
                    id=int(row[index["id"]]),
                    submission_id=int(row[index["submission_id"]]),
                    question_id=int(row[index["question_id"]]),
                    answer_value=int(row[index["answer_value"]]),
                )
                answers_by_submission[answer.submission_id].append(answer)

    return users, profiles, assessments, submissions, reports_by_submission, questions, answers_by_submission


def submission_sort_key(submission: Submission) -> tuple:
    try:
        dt = datetime.fromisoformat(submission.submitted_at.replace(" ", "T"))
    except Exception:
        dt = datetime.min
    return (dt, submission.id)


def export_rows(sql_path: Path):
    users, profiles, assessments, submissions, reports_by_submission, questions, answers_by_submission = load_backup(sql_path)

    submissions_by_user: Dict[int, List[Submission]] = defaultdict(list)
    for submission in submissions.values():
        submissions_by_user[submission.user_id].append(submission)
    for item in submissions_by_user.values():
        item.sort(key=submission_sort_key, reverse=True)

    latest_by_user_assessment: Dict[tuple[int, str], Submission] = {}
    for user_id, user_subs in submissions_by_user.items():
        for submission in user_subs:
            paper_id = assessments[submission.assessment_id].test_paper_id
            key = (user_id, paper_id)
            if key not in latest_by_user_assessment:
                latest_by_user_assessment[key] = submission

    master_rows: List[List[object]] = []
    report_rows: List[List[object]] = []
    answer_rows: List[List[object]] = []

    user_seq = 1
    for user in sorted(users.values(), key=lambda item: item.id):
        if user.role != "USER":
            continue
        profile = profiles.get(user.id)
        scl = latest_by_user_assessment.get((user.id, "SCL90-2078"))
        sds = latest_by_user_assessment.get((user.id, "SDS-2082"))
        sas = latest_by_user_assessment.get((user.id, "SAS-2083"))
        composite = build_composite_result(scl, sds, sas)
        master_rows.append([
            user_seq,
            user.id,
            profile.name if profile else "",
            user.email,
            profile.gender if profile else "",
            profile.id_card if profile else "",
            profile.phone if profile else "",
            profile.birth_month if profile else "",
            profile.education if profile else "",
            profile.school_major if profile else "",
            scl.raw_score if scl else "",
            scl.level if scl else "",
            sds.raw_score if sds else "",
            sds.level if sds else "",
            sas.raw_score if sas else "",
            sas.level if sas else "",
            composite.score if composite.score is not None else "",
            composite.level,
            composite.progress,
        ])
        user_seq += 1

    for submission in sorted(submissions.values(), key=submission_sort_key, reverse=True):
        user = users.get(submission.user_id)
        if user is None or user.role != "USER":
            continue
        profile = profiles.get(user.id)
        assessment = assessments[submission.assessment_id]
        report = reports_by_submission.get(submission.id)
        label = ASSESSMENT_LABELS.get(assessment.test_paper_id, assessment.test_paper_id)
        report_rows.append([
            user.id,
            profile.name if profile else "",
            user.email,
            label,
            submission.raw_score,
            submission.level,
            report.summary if report else "",
            report.suggestions if report else "",
            submission.submitted_at,
            submission.id,
            report.id if report else "",
        ])

        question_list = [
            (questions[answer.question_id], answer)
            for answer in answers_by_submission.get(submission.id, [])
            if answer.question_id in questions
        ]
        question_list.sort(key=lambda item: (item[0].sort_order, item[0].id, item[1].id))
        for question, answer in question_list:
            answer_rows.append([
                user.id,
                profile.name if profile else "",
                user.email,
                label,
                submission.id,
                question.code,
                question.dimension,
                question.sort_order,
                question.prompt,
                answer.answer_value,
                submission.submitted_at,
            ])

    return master_rows, report_rows, answer_rows


def write_csv(path: Path, header: List[str], rows: List[List[object]]) -> None:
    with path.open("w", newline="", encoding="utf-8-sig") as handle:
        writer = csv.writer(handle)
        writer.writerow(header)
        writer.writerows(rows)


def write_xlsx(path: Path, datasets: List[tuple[str, List[str], List[List[object]]]]) -> None:
    try:
        from openpyxl import Workbook
        from openpyxl.styles import Alignment, Font, PatternFill
    except Exception as exc:
        raise RuntimeError("当前 Python 缺少 openpyxl，请改用 Codex 自带 Python 运行。") from exc

    workbook = Workbook()
    workbook.remove(workbook.active)
    header_fill = PatternFill("solid", fgColor="FFF2CC")
    header_font = Font(bold=True)
    wrap = Alignment(vertical="top", wrap_text=True)

    for sheet_name, header, rows in datasets:
        sheet = workbook.create_sheet(title=sheet_name[:31])
        sheet.append(header)
        for col in sheet[1]:
            col.fill = header_fill
            col.font = header_font
            col.alignment = wrap
        for row in rows:
            sheet.append(row)
        for row in sheet.iter_rows():
            for cell in row:
                cell.alignment = wrap
        for column_cells in sheet.columns:
            length = max(len(str(cell.value or "")) for cell in column_cells)
            sheet.column_dimensions[column_cells[0].column_letter].width = min(max(length + 2, 10), 48)
        sheet.freeze_panes = "A2"
        sheet.auto_filter.ref = sheet.dimensions

    workbook.save(path)


def main() -> int:
    parser = argparse.ArgumentParser(description="从 davos_assessment_full.sql 离线导出用户测评台账")
    parser.add_argument("--input", required=True, help="SQL 备份文件路径，例如 /Users/.../davos_assessment_full.sql")
    parser.add_argument("--output-dir", default="", help="导出目录，默认在输入文件同目录下生成 export-时间戳")
    args = parser.parse_args()

    sql_path = Path(args.input).expanduser().resolve()
    if not sql_path.exists():
        print(f"找不到输入文件：{sql_path}", file=sys.stderr)
        return 1

    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    output_dir = Path(args.output_dir).expanduser().resolve() if args.output_dir else sql_path.parent / f"davos-export-{stamp}"
    output_dir.mkdir(parents=True, exist_ok=True)

    master_rows, report_rows, answer_rows = export_rows(sql_path)

    master_header = [
        "序号", "用户ID", "姓名", "邮箱", "性别", "身份证号", "联系方式", "出生日期", "学历", "院校/专业",
        "SCL90原始分", "SCL90等级", "SDS标准分", "SDS等级", "SAS标准分", "SAS等级",
        "综合分", "综合风险等级", "完成进度",
    ]
    report_header = [
        "用户ID", "姓名", "邮箱", "量表", "分数", "等级", "报告结论", "报告建议", "完成时间", "SubmissionID", "ReportID",
    ]
    answer_header = [
        "用户ID", "姓名", "邮箱", "量表", "SubmissionID", "题号", "维度", "排序", "题目", "答案", "完成时间",
    ]

    write_csv(output_dir / "01_用户总览.csv", master_header, master_rows)
    write_csv(output_dir / "02_报告摘要.csv", report_header, report_rows)
    write_csv(output_dir / "03_逐题答案.csv", answer_header, answer_rows)
    write_xlsx(
        output_dir / "达沃斯心理测评导出.xlsx",
        [
            ("用户总览", master_header, master_rows),
            ("报告摘要", report_header, report_rows),
            ("逐题答案", answer_header, answer_rows),
        ],
    )

    print("导出完成：")
    print(output_dir)
    print(f"用户总览：{len(master_rows)} 行")
    print(f"报告摘要：{len(report_rows)} 行")
    print(f"逐题答案：{len(answer_rows)} 行")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
