"""LLM-assisted epoch configuration wizard.

This module provides tools to infer Statistical Learning paradigm parameters
from event markers using both heuristic analysis and LLM assistance.

Supports:
- OpenCode server (opencode serve)
- Direct Anthropic API
- Any OpenAI-compatible endpoint (Ollama, etc.)

Example:
    >>> from autocleaneeg_sl.wizard import infer_config, run_wizard
    >>> events = [11, 2, 4, 8, 3, 5, 9, 7, 6, 1, 12, 10] * 200
    >>> config = infer_config(events)  # Heuristic only
    >>> config = run_wizard(events, provider="opencode")  # With LLM
"""

from collections import Counter
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Literal
import re

import numpy as np

try:
    import requests
    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False

try:
    import yaml
    HAS_YAML = True
except ImportError:
    HAS_YAML = False


# Default OpenCode server URL
OPENCODE_URL = "http://127.0.0.1:4096"

# Provider types
ProviderType = Literal["opencode", "anthropic", "openai", "ollama"]


@dataclass
class WizardResult:
    """Result from the configuration wizard."""

    # Inferred parameters
    syllable_codes: list[int]
    word_onset_codes: list[int]
    words: list[list[int]]
    soa_ms: float
    syllable_rate_hz: float
    word_rate_hz: float

    # Metadata
    confidence: float
    method: str  # "heuristic" or "llm"
    notes: list[str] = field(default_factory=list)
    questions: list[str] = field(default_factory=list)
    raw_response: Optional[str] = None
    validated_yaml: Optional[str] = None  # LLM-generated YAML if validated

    def to_yaml(self) -> str:
        """Generate YAML config string for autocleaneeg-sl.

        If the LLM provided a validated YAML, use it directly.
        Otherwise, generate from parsed values.
        """
        # Use LLM's validated YAML if available
        if self.validated_yaml:
            header = f"# Auto-generated configuration\n# Method: {self.method} (confidence: {self.confidence:.0%})\n\n"
            return header + self.validated_yaml

        # Otherwise generate from parsed values
        lines = [
            f"# Auto-generated configuration",
            f"# Method: {self.method} (confidence: {self.confidence:.0%})",
            "",
            "paradigm:",
            f"  name: inferred-sl-paradigm",
            f"  syllable_soa_ms: {self.soa_ms}",
            f"  syllables_per_epoch: 36",
            f"  jitter_threshold_ms: {self.soa_ms + 20}",
            "",
            f"syllable_codes: {self.syllable_codes}",
            "",
            "languages:",
            "  stream1:",
            f"    word_onset_codes: {self.word_onset_codes}",
            "    words:",
        ]

        for word in self.words:
            lines.append(f"      - \"{'-'.join(map(str, word))}\"")

        lines.extend([
            "",
            "analysis:",
            "  freq_range: [0.5, 10.0]",
            "  target_freqs:",
            f"    word: {self.word_rate_hz:.3f}",
            f"    syllable: {self.syllable_rate_hz:.3f}",
        ])

        if self.notes:
            lines.extend(["", "# Notes:"])
            for note in self.notes:
                lines.append(f"# - {note}")

        if self.questions:
            lines.extend(["", "# Questions to verify:"])
            for q in self.questions:
                lines.append(f"# - {q}")

        return "\n".join(lines)

    def to_config(self):
        """Convert to SLConfig object."""
        from .config import (
            SLConfig, ParadigmConfig, LanguageConfig, AnalysisConfig
        )

        return SLConfig(
            paradigm=ParadigmConfig(
                name="inferred-sl-paradigm",
                syllable_soa_ms=self.soa_ms,
                syllables_per_epoch=36,
                jitter_threshold_ms=self.soa_ms + 20,
            ),
            syllable_codes=self.syllable_codes,
            languages={
                "stream1": LanguageConfig(
                    word_onset_codes=self.word_onset_codes,
                ),
            },
            subjects={},
            analysis=AnalysisConfig(
                freq_range=(0.5, 10.0),
                target_freqs={
                    "word": self.word_rate_hz,
                    "syllable": self.syllable_rate_hz,
                },
            ),
        )


# =============================================================================
# Heuristic Inference (no LLM)
# =============================================================================


def detect_triplet_patterns(events: list[int]) -> tuple[list[list[int]], float]:
    """Detect word (triplet) patterns from event sequence.

    Analyzes the event sequence to find consistent triplet groupings
    that indicate word boundaries in SL paradigms.

    Args:
        events: List of event codes.

    Returns:
        Tuple of (words, confidence) where words is a list of triplets.
    """
    # Count triplets at word boundaries (every 3rd position)
    triplet_counts = Counter()
    for i in range(0, len(events) - 2, 3):
        triplet = tuple(events[i:i + 3])
        triplet_counts[triplet] += 1

    # Get the most common triplets
    unique_codes = set(events)
    n_syllables = len(unique_codes)

    # Expect n_syllables / 3 words for triplet structure
    expected_words = n_syllables // 3 if n_syllables >= 3 else 1

    # Get top triplets that cover all syllables
    words = []
    used_codes = set()

    for triplet, count in triplet_counts.most_common():
        triplet_codes = set(triplet)
        # Only add if it uses new codes
        if not triplet_codes & used_codes:
            words.append(list(triplet))
            used_codes.update(triplet_codes)

        if len(words) >= expected_words:
            break

    # Calculate confidence based on coverage and consistency
    if not words:
        return [], 0.0

    coverage = len(used_codes) / len(unique_codes) if unique_codes else 0
    consistency = min(1.0, len(words) / expected_words) if expected_words else 0
    confidence = (coverage + consistency) / 2

    return words, confidence


def estimate_soa(
    event_samples: np.ndarray,
    sfreq: float,
) -> tuple[float, float]:
    """Estimate SOA from event timing.

    Args:
        event_samples: Sample indices of events.
        sfreq: Sampling frequency in Hz.

    Returns:
        Tuple of (soa_ms, confidence).
    """
    if len(event_samples) < 2:
        return 300.0, 0.0

    intervals_ms = np.diff(event_samples) / sfreq * 1000

    # Filter reasonable intervals (100-500ms typical for SL)
    valid = intervals_ms[(intervals_ms > 100) & (intervals_ms < 500)]

    if len(valid) == 0:
        return 300.0, 0.2

    soa = float(np.median(valid))
    std = float(np.std(valid))
    cv = std / soa if soa > 0 else 1.0
    confidence = max(0.3, 1.0 - cv * 2)

    return soa, confidence


def infer_config(
    events: list[int],
    event_samples: Optional[np.ndarray] = None,
    sfreq: Optional[float] = None,
    soa_ms: Optional[float] = None,
) -> WizardResult:
    """Infer SL paradigm config using heuristics only.

    This is fast and doesn't require an LLM, but may be less accurate
    for complex or unusual paradigms.

    Args:
        events: List of event codes.
        event_samples: Optional sample indices for timing estimation.
        sfreq: Sampling frequency (required if event_samples provided).
        soa_ms: Known SOA in ms (if not inferring from timing).

    Returns:
        WizardResult with inferred configuration.
    """
    unique_codes = sorted(set(events))
    notes = []
    questions = []

    # Detect word structure
    words, word_confidence = detect_triplet_patterns(events)

    if words:
        word_onset_codes = [w[0] for w in words]
        syllables_per_word = len(words[0])
        notes.append(f"Detected {len(words)} words with {syllables_per_word} syllables each")
    else:
        word_onset_codes = unique_codes[:4] if len(unique_codes) >= 4 else unique_codes
        words = [[c] for c in word_onset_codes]
        syllables_per_word = 3
        questions.append("Could not detect word structure - please specify word onset codes")

    # Estimate SOA
    soa_confidence = 0.0
    if soa_ms is None:
        if event_samples is not None and sfreq is not None:
            soa_ms, soa_confidence = estimate_soa(event_samples, sfreq)
            notes.append(f"Estimated SOA: {soa_ms:.1f}ms")
        else:
            soa_ms = 300.0
            notes.append("Using default SOA of 300ms")
            questions.append("What is the SOA (time between syllables)?")

    # Calculate rates
    syllable_rate = 1000 / soa_ms
    word_rate = syllable_rate / syllables_per_word

    # Overall confidence
    confidence = (word_confidence + soa_confidence) / 2 if soa_confidence > 0 else word_confidence * 0.7

    return WizardResult(
        syllable_codes=unique_codes,
        word_onset_codes=word_onset_codes,
        words=words,
        soa_ms=soa_ms,
        syllable_rate_hz=syllable_rate,
        word_rate_hz=word_rate,
        confidence=confidence,
        method="heuristic",
        notes=notes,
        questions=questions,
    )


# =============================================================================
# LLM-Assisted Inference
# =============================================================================


def _build_llm_prompt(
    events: list[int],
    soa_ms: Optional[float] = None,
    syllable_map: Optional[dict[int, str]] = None,
    description: Optional[str] = None,
) -> str:
    """Build prompt for LLM analysis."""
    events_sample = " ".join(map(str, events[:200]))
    unique_codes = sorted(set(events))

    prompt = f"""Analyze this Statistical Learning EEG paradigm and identify the word structure.

## Event Data
- Total events: {len(events)}
- Unique codes: {unique_codes}
- First 200 events: {events_sample}...

"""

    if syllable_map:
        prompt += "## Syllable Mapping\n"
        for code, name in sorted(syllable_map.items()):
            prompt += f"- Code {code} = '{name}'\n"
        prompt += "\n"

    if soa_ms:
        prompt += f"## Timing\n- SOA: {soa_ms}ms\n- Syllable rate: {1000/soa_ms:.2f} Hz\n\n"

    if description:
        prompt += f"## Additional Info\n{description}\n\n"

    prompt += """## Task
1. Analyze triplet patterns to identify the words (usually 3-4 words of 3 syllables each)
2. Identify word onset codes (first syllable of each word)
3. Generate YAML config with these exact fields:

```yaml
paradigm:
  name: string
  syllable_soa_ms: float
  syllables_per_epoch: 36
  jitter_threshold_ms: float

syllable_codes: [list of ints]

languages:
  stream1:
    word_onset_codes: [list of ints]
    words:
      - "code-code-code"

analysis:
  freq_range: [0.5, 10.0]
  target_freqs:
    word: float
    syllable: float
```

Show brief analysis then the YAML."""

    return prompt


def _validate_yaml_schema(config: dict) -> tuple[bool, list[str]]:
    """Validate LLM-generated config against our schema.

    Returns:
        Tuple of (is_valid, list of issues).
    """
    issues = []

    # Required top-level keys
    if "paradigm" not in config:
        issues.append("Missing 'paradigm' section")
    else:
        paradigm = config["paradigm"]
        if "syllable_soa_ms" not in paradigm:
            issues.append("Missing 'paradigm.syllable_soa_ms'")

    if "syllable_codes" not in config:
        issues.append("Missing 'syllable_codes'")

    if "languages" not in config:
        issues.append("Missing 'languages' section")
    else:
        if "stream1" not in config["languages"]:
            issues.append("Missing 'languages.stream1'")
        else:
            stream = config["languages"]["stream1"]
            if "word_onset_codes" not in stream:
                issues.append("Missing 'languages.stream1.word_onset_codes'")

    return len(issues) == 0, issues


def _parse_llm_response(response: str, events: list[int], soa_ms: float) -> WizardResult:
    """Parse LLM response into WizardResult.

    If the LLM's YAML passes validation, it's stored directly for output.
    """
    unique_codes = sorted(set(events))

    # Extract YAML from response
    yaml_match = re.search(r'```ya?ml\n(.*?)```', response, re.DOTALL)

    words = []
    word_onset_codes = []
    validated_yaml = None

    if yaml_match and HAS_YAML:
        yaml_str = yaml_match.group(1)
        try:
            config = yaml.safe_load(yaml_str)

            # Validate schema
            is_valid, issues = _validate_yaml_schema(config)

            if is_valid:
                # Store validated YAML for direct output
                validated_yaml = yaml_str.strip()

            # Extract word onset codes
            if "languages" in config:
                stream = config["languages"].get("stream1", {})
                word_onset_codes = stream.get("word_onset_codes", [])

                # Extract words
                for word_str in stream.get("words", []):
                    if isinstance(word_str, str):
                        # Parse "11-2-4" format
                        word = [int(x) for x in re.findall(r'\d+', word_str)]
                        if word:
                            words.append(word)
                    elif isinstance(word_str, list):
                        words.append(word_str)

            # Get SOA from config if present
            if "paradigm" in config:
                soa_ms = config["paradigm"].get("syllable_soa_ms", soa_ms)

        except Exception:
            pass

    # Fallback: try to extract from text
    if not word_onset_codes:
        # Look for patterns like "word onset codes: [1, 8, 9, 11]"
        onset_match = re.search(r'word.?onset.?codes?[:\s]+\[([^\]]+)\]', response, re.IGNORECASE)
        if onset_match:
            word_onset_codes = [int(x) for x in re.findall(r'\d+', onset_match.group(1))]

    # If still no words, infer from onset codes
    if not words and word_onset_codes:
        # Create placeholder words
        words = [[c, c, c] for c in word_onset_codes]  # Placeholder

    # Calculate rates
    syllable_rate = 1000 / soa_ms
    syllables_per_word = len(words[0]) if words else 3
    word_rate = syllable_rate / syllables_per_word

    return WizardResult(
        syllable_codes=unique_codes,
        word_onset_codes=word_onset_codes or unique_codes[:4],
        words=words or [[c] for c in unique_codes[:4]],
        soa_ms=soa_ms,
        syllable_rate_hz=syllable_rate,
        word_rate_hz=word_rate,
        confidence=0.95 if validated_yaml else (0.8 if word_onset_codes else 0.5),
        method="llm",
        notes=["Configuration generated by LLM analysis"] + (["YAML validated against schema"] if validated_yaml else []),
        raw_response=response,
        validated_yaml=validated_yaml,
    )


def _call_opencode(prompt: str, base_url: str = OPENCODE_URL) -> Optional[str]:
    """Call OpenCode server API."""
    if not HAS_REQUESTS:
        raise ImportError("requests package required: pip install requests")

    # Check health
    try:
        health = requests.get(f"{base_url}/global/health", timeout=5)
        if not health.ok:
            raise ConnectionError("OpenCode server not responding")
    except requests.exceptions.ConnectionError:
        raise ConnectionError(
            f"Cannot connect to OpenCode server at {base_url}\n"
            "Start it with: opencode serve"
        )

    # Create session
    session_resp = requests.post(
        f"{base_url}/session",
        json={},
        headers={"Content-Type": "application/json"},
    )
    if not session_resp.ok:
        raise RuntimeError("Failed to create OpenCode session")

    session_id = session_resp.json()["id"]

    # Send prompt
    msg_resp = requests.post(
        f"{base_url}/session/{session_id}/message",
        json={"parts": [{"type": "text", "text": prompt}]},
        headers={"Content-Type": "application/json"},
        timeout=120,
    )

    if msg_resp.ok:
        data = msg_resp.json()
        text_parts = [
            p["text"] for p in data.get("parts", [])
            if p.get("type") == "text"
        ]
        return "\n".join(text_parts)

    return None


def _call_anthropic(prompt: str, api_key: Optional[str] = None) -> Optional[str]:
    """Call Anthropic API directly."""
    if not HAS_REQUESTS:
        raise ImportError("requests package required: pip install requests")

    import os
    api_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        raise ValueError("ANTHROPIC_API_KEY not set")

    resp = requests.post(
        "https://api.anthropic.com/v1/messages",
        headers={
            "x-api-key": api_key,
            "anthropic-version": "2023-06-01",
            "content-type": "application/json",
        },
        json={
            "model": "claude-sonnet-4-20250514",
            "max_tokens": 2000,
            "messages": [{"role": "user", "content": prompt}],
        },
        timeout=120,
    )

    if resp.ok:
        data = resp.json()
        return data["content"][0]["text"]
    return None


def _call_openai_compatible(
    prompt: str,
    base_url: str,
    api_key: str = "not-needed",
    model: str = "llama3",
) -> Optional[str]:
    """Call OpenAI-compatible API (Ollama, etc.)."""
    if not HAS_REQUESTS:
        raise ImportError("requests package required: pip install requests")

    resp = requests.post(
        f"{base_url}/chat/completions",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        json={
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": 2000,
        },
        timeout=120,
    )

    if resp.ok:
        data = resp.json()
        return data["choices"][0]["message"]["content"]
    return None


def run_wizard(
    events: list[int],
    provider: ProviderType = "opencode",
    soa_ms: Optional[float] = None,
    syllable_map: Optional[dict[int, str]] = None,
    description: Optional[str] = None,
    base_url: Optional[str] = None,
    api_key: Optional[str] = None,
    model: Optional[str] = None,
    verbose: bool = False,
) -> WizardResult:
    """Run the LLM-assisted configuration wizard.

    Args:
        events: List of event codes from EEG/stimulus file.
        provider: LLM provider ("opencode", "anthropic", "openai", "ollama").
        soa_ms: Known SOA in milliseconds.
        syllable_map: Optional mapping of codes to syllable names.
        description: Optional additional paradigm description.
        base_url: Custom API base URL (for ollama/openai).
        api_key: API key (for anthropic/openai).
        model: Model name (for ollama/openai).
        verbose: Print progress.

    Returns:
        WizardResult with LLM-inferred configuration.

    Example:
        >>> events = load_events_from_file("StructuredOrder.txt")
        >>> result = run_wizard(events, provider="opencode", soa_ms=300)
        >>> print(result.to_yaml())
    """
    # Default SOA
    if soa_ms is None:
        soa_ms = 300.0

    # Build prompt
    prompt = _build_llm_prompt(events, soa_ms, syllable_map, description)

    if verbose:
        print(f"Calling {provider} API...")

    # Call appropriate provider
    if provider == "opencode":
        url = base_url or OPENCODE_URL
        response = _call_opencode(prompt, url)
    elif provider == "anthropic":
        response = _call_anthropic(prompt, api_key)
    elif provider in ("openai", "ollama"):
        url = base_url or ("http://localhost:11434/v1" if provider == "ollama" else "https://api.openai.com/v1")
        response = _call_openai_compatible(prompt, url, api_key or "not-needed", model or "llama3")
    else:
        raise ValueError(f"Unknown provider: {provider}")

    if response is None:
        if verbose:
            print("LLM call failed, falling back to heuristics")
        return infer_config(events, soa_ms=soa_ms)

    if verbose:
        print("Parsing LLM response...")

    return _parse_llm_response(response, events, soa_ms)


# =============================================================================
# File Loading Utilities
# =============================================================================


def load_events_from_file(filepath: str | Path) -> list[int]:
    """Load event codes from a text file.

    Supports space-separated, comma-separated, or newline-separated integers.

    Args:
        filepath: Path to events file.

    Returns:
        List of event codes.
    """
    filepath = Path(filepath)
    content = filepath.read_text()
    return [int(x) for x in re.findall(r'\d+', content)]


def load_events_from_eeg(filepath: str | Path) -> tuple[list[int], np.ndarray, float]:
    """Load events from an EEG file (.set).

    Args:
        filepath: Path to .set file.

    Returns:
        Tuple of (event_codes, event_samples, sfreq).
    """
    import mne

    filepath = Path(filepath)

    # Try loading as raw first
    try:
        raw = mne.io.read_raw_eeglab(str(filepath), preload=False, verbose=False)
        events, _ = mne.events_from_annotations(raw, verbose=False)
        sfreq = raw.info["sfreq"]
    except Exception:
        # Try as epochs
        epochs = mne.read_epochs_eeglab(str(filepath), verbose=False)
        events = epochs.events
        sfreq = epochs.info["sfreq"]

    codes = events[:, 2].tolist()
    samples = events[:, 0]

    return codes, samples, sfreq
