"""Extract training examples from Obsidian vault.

Looks for reflection patterns, decision analyses, and personal insights
to create training data that teaches the model to reason like San.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator

from companion.config import Config
from companion.rag.chunker import chunk_file


@dataclass
class TrainingExample:
    """A single training example in conversation format."""

    messages: list[dict[str, str]]
    source_file: str
    tags: list[str]
    date: str | None = None

    def to_dict(self) -> dict:
        return {
            "messages": self.messages,
            "source_file": self.source_file,
            "tags": self.tags,
            "date": self.date,
        }


# Tags that indicate reflection/decision content
REFLECTION_TAGS = {
    "#reflection",
    "#decision",
    "#learning",
    "#insight",
    "#analysis",
    "#pondering",
    "#evaluation",
    "#takeaway",
    "#realization",
}

# Patterns that suggest reflection content
REFLECTION_PATTERNS = [
    r"(?i)I think\s+",
    r"(?i)I feel\s+",
    r"(?i)I realize\s+",
    r"(?i)I wonder\s+",
    r"(?i)I should\s+",
    r"(?i)I need to\s+",
    r"(?i)The reason\s+",
    r"(?i)What if\s+",
    r"(?i)Maybe\s+",
    r"(?i)Perhaps\s+",
    r"(?i)On one hand.*?on the other hand",
    r"(?i)Pros?:.*?Cons?:",
    r"(?i)Weighing.*?(?:options?|choices?|alternatives?)",
    r"(?i)Ultimately.*?decided",
    r"(?i)Looking back",
    r"(?i)In hindsight",
    r"(?i)I've learned\s+",
    r"(?i)The lesson\s+",
]


def _has_reflection_tags(text: str) -> bool:
    """Check if text contains reflection-related hashtags."""
    hashtags = set(re.findall(r"#\w+", text))
    return bool(hashtags & REFLECTION_TAGS)


def _has_reflection_patterns(text: str) -> bool:
    """Check if text contains reflection language patterns."""
    for pattern in REFLECTION_PATTERNS:
        if re.search(pattern, text):
            return True
    return False


def _is_likely_reflection(text: str) -> bool:
    """Determine if a text chunk is likely a reflection."""
    # Must have reflection tags OR strong reflection patterns
    has_tags = _has_reflection_tags(text)
    has_patterns = _has_reflection_patterns(text)

    # Require at least one indicator
    return has_tags or has_patterns


def _extract_date_from_filename(filename: str) -> str | None:
    """Extract date from filename patterns like 2026-04-12 or "12 Apr 2026"."""
    # ISO format: 2026-04-12
    m = re.search(r"(\d{4}-\d{2}-\d{2})", filename)
    if m:
        return m.group(1)

    # Human format: 12-Apr-2026 or "12 Apr 2026"
    m = re.search(r"(\d{1,2}[-\s][A-Za-z]{3}[-\s]\d{4})", filename)
    if m:
        return m.group(1).replace(" ", "-").replace("--", "-")

    return None


def _create_training_prompt(chunk_text: str) -> str:
    """Create a user prompt that elicits a reflection similar to the example."""
    # Extract context from the chunk to form a relevant question

    # Look for decision patterns
    if re.search(r"(?i)decided|decision|chose|choice", chunk_text):
        return "I'm facing a decision. How should I think through this?"

    # Look for relationship content
    if re.search(r"(?i)friend|relationship|person|people|someone", chunk_text):
        return "What do you think about this situation with someone I'm close to?"

    # Look for work/career content
    if re.search(r"(?i)work|job|career|project|professional", chunk_text):
        return "I'm thinking about something at work. What's your perspective?"

    # Look for health content
    if re.search(r"(?i)health|mental|physical|stress|wellness", chunk_text):
        return "I've been thinking about my well-being. What do you notice?"

    # Look for financial content
    if re.search(r"(?i)money|finance|financial|invest|saving|spending", chunk_text):
        return "I'm considering a financial decision. How should I evaluate this?"

    # Default prompt
    return "I'm reflecting on something. What patterns do you see?"


def _create_training_example(
    chunk_text: str, source_file: str, tags: list[str], date: str | None
) -> TrainingExample | None:
    """Convert a chunk into a training example if it meets criteria."""

    # Skip if too short
    if len(chunk_text) < 100:
        return None

    # Skip if not a reflection
    if not _is_likely_reflection(chunk_text):
        return None

    # Clean up the text
    cleaned = chunk_text.strip()

    # Create prompt-response pair
    prompt = _create_training_prompt(cleaned)

    # The response is your reflection/insight
    messages = [
        {"role": "system", "content": "You are a thoughtful, reflective companion."},
        {"role": "user", "content": prompt},
        {"role": "assistant", "content": cleaned},
    ]

    return TrainingExample(
        messages=messages,
        source_file=source_file,
        tags=tags,
        date=date,
    )


class TrainingDataExtractor:
    """Extracts training examples from Obsidian vault."""

    def __init__(self, config: Config):
        self.config = config
        self.vault_path = Path(config.vault.path)
        self.examples: list[TrainingExample] = []

    def extract(self) -> list[TrainingExample]:
        """Extract all training examples from the vault."""
        self.examples = []

        # Walk vault for markdown files
        for md_file in self.vault_path.rglob("*.md"):
            # Skip denied directories
            relative = md_file.relative_to(self.vault_path)
            if any(part.startswith(".") for part in relative.parts):
                continue
            if any(
                part in [".obsidian", ".trash", "zzz-Archive"]
                for part in relative.parts
            ):
                continue

            # Extract from this file
            file_examples = self._extract_from_file(md_file)
            self.examples.extend(file_examples)

        return self.examples

    def _extract_from_file(self, md_file: Path) -> list[TrainingExample]:
        """Extract training examples from a single file."""
        examples = []

        try:
            text = md_file.read_text(encoding="utf-8")
        except Exception:
            return examples

        # Get date from filename
        date = _extract_date_from_filename(md_file.name)

        # Split into sections (by headers like #Section or # Tag:)
        # Pattern matches lines starting with # that have content
        sections = re.split(r"\n(?=#[^#])", text)

        for section in sections:
            section = section.strip()
            if not section or len(section) < 50:  # Skip very short sections
                continue

            # Extract tags from section
            hashtags = re.findall(r"#[\w\-]+", section)

            # Try to create training example
            example = _create_training_example(
                chunk_text=section,
                source_file=str(md_file.relative_to(self.vault_path)).replace(
                    "\\", "/"
                ),
                tags=hashtags,
                date=date,
            )

            if example:
                examples.append(example)

        # If no reflection sections found, try the whole file
        if not examples and len(text) >= 100:
            hashtags = re.findall(r"#[\w\-]+", text)
            if _is_likely_reflection(text):
                example = _create_training_example(
                    chunk_text=text.strip(),
                    source_file=str(md_file.relative_to(self.vault_path)).replace(
                        "\\", "/"
                    ),
                    tags=hashtags,
                    date=date,
                )
                if example:
                    examples.append(example)

        return examples

        # Get date from filename
        date = _extract_date_from_filename(md_file.name)

        # Split into sections (by headers)
        sections = re.split(r"\n(?=#+\s)", text)

        for section in sections:
            section = section.strip()
            if not section:
                continue

            # Extract tags from section
            hashtags = re.findall(r"#\w+", section)

            # Try to create training example
            example = _create_training_example(
                chunk_text=section,
                source_file=str(md_file.relative_to(self.vault_path)).replace(
                    "\\", "/"
                ),
                tags=hashtags,
                date=date,
            )

            if example:
                examples.append(example)

        return examples

    def save_to_jsonl(self, output_path: Path) -> int:
        """Save extracted examples to JSONL file."""
        with open(output_path, "w", encoding="utf-8") as f:
            for example in self.examples:
                f.write(json.dumps(example.to_dict(), ensure_ascii=False) + "\n")
        return len(self.examples)

    def get_stats(self) -> dict:
        """Get statistics about extracted examples."""
        if not self.examples:
            return {"total": 0}

        tag_counts: dict[str, int] = {}
        for ex in self.examples:
            for tag in ex.tags:
                tag_counts[tag] = tag_counts.get(tag, 0) + 1

        return {
            "total": len(self.examples),
            "avg_length": sum(len(ex.messages[2]["content"]) for ex in self.examples)
            // len(self.examples),
            "top_tags": sorted(tag_counts.items(), key=lambda x: x[1], reverse=True)[
                :10
            ],
        }


def extract_training_data(
    config_path: str = "config.json",
) -> tuple[list[TrainingExample], dict]:
    """Convenience function to extract training data from vault.

    Returns:
        Tuple of (examples list, stats dict)
    """
    from companion.config import load_config

    config = load_config(config_path)
    extractor = TrainingDataExtractor(config)
    examples = extractor.extract()
    stats = extractor.get_stats()

    return examples, stats
