"""
Annotation saving utilities for IVUS annotation tool.
Handles CSV file creation, saving annotations, and tracking progress.
"""

import os
import csv
import pandas as pd
from datetime import datetime
from typing import List, Set, Union, Optional


def initialize_csv(csv_path: str) -> None:
    """
    Create annotation CSV with headers if it doesn't exist.

    Args:
        csv_path: Path to CSV file

    CSV Columns:
        - timestamp: When annotation was saved (YYYY-MM-DD HH:MM:SS)
        - case_id: Case number (int or float)
        - prediction: Complication prediction ("あり" or "なし")
        - confidence: Confidence level (0-100)
        - reasons: Selected reasons (semicolon-separated)
        - comment: Free text comment
        - annotator: Name of annotator
        - ground_truth: Ground truth label from Excel (True/False/None)
    """
    if not os.path.exists(csv_path):
        with open(csv_path, 'w', encoding='utf-8-sig', newline='') as f:
            writer = csv.writer(f)
            writer.writerow([
                'timestamp',
                'case_id',
                'prediction',
                'confidence',
                'reasons',
                'comment',
                'annotator',
                'ground_truth'
            ])


def save_annotation(
    csv_path: str,
    case_id: Union[int, float],
    prediction: str,
    confidence: int,
    reasons: List[str],
    comment: str,
    annotator: str,
    ground_truth: Optional[bool]
) -> None:
    """
    Save annotation to CSV file in append mode.

    Args:
        csv_path: Path to CSV file
        case_id: Case number
        prediction: "あり" or "なし"
        confidence: 0-100
        reasons: List of selected reasons
        comment: Free text comment
        annotator: Annotator name
        ground_truth: Ground truth label (True/False/None if not available)

    Example:
        >>> save_annotation(
        ...     "/path/to/annotations_tanaka.csv",
        ...     134,
        ...     "あり",
        ...     75,
        ...     ["石灰化プラークが多い", "減衰プラークが多い"],
        ...     "明確な所見あり",
        ...     "tanaka",
        ...     True
        ... )
    """
    initialize_csv(csv_path)

    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    reasons_str = "; ".join(reasons)  # Join multiple reasons with semicolon

    # Convert ground_truth to string representation
    if ground_truth is None:
        gt_str = ""
    elif ground_truth:
        gt_str = "True"
    else:
        gt_str = "False"

    with open(csv_path, 'a', encoding='utf-8-sig', newline='') as f:
        writer = csv.writer(f)
        writer.writerow([
            timestamp,
            case_id,
            prediction,
            confidence,
            reasons_str,
            comment,
            annotator,
            gt_str
        ])


def get_annotated_cases(csv_path: str, annotator: str = None) -> Set[Union[int, float]]:
    """
    Get set of case IDs that have been annotated.

    Args:
        csv_path: Path to CSV file
        annotator: Optional filter by specific annotator

    Returns:
        Set of case IDs that have annotations

    Example:
        >>> get_annotated_cases("/path/to/annotations_tanaka.csv", "tanaka")
        {134, 134.1, 135, 136, ...}
    """
    if not os.path.exists(csv_path):
        return set()

    try:
        df = pd.read_csv(csv_path, encoding='utf-8-sig')

        if annotator:
            df = df[df['annotator'] == annotator]

        # Convert case_ids to appropriate type (int or float)
        case_ids = set()
        for case_id in df['case_id'].unique():
            try:
                # Try to convert to float first
                case_id_float = float(case_id)
                # If it's a whole number, store as int, otherwise as float
                if case_id_float.is_integer():
                    case_ids.add(int(case_id_float))
                else:
                    case_ids.add(case_id_float)
            except (ValueError, TypeError):
                continue

        return case_ids
    except Exception as e:
        print(f"Warning: Could not read annotations from {csv_path}: {e}")
        return set()