from __future__ import annotations

import json
import logging
import re
import shutil
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from fastapi import APIRouter, HTTPException, Header

from prompt_auto_lab.api.models import (
    ImportSessionRequest,
    Session,
    SessionStatus,
    MetricsSnapshot,
    AgentEvent,
    AgentEventType,
)
from prompt_auto_lab.api.services import user_service
from prompt_auto_lab.core.store_dir import get_store_dir

router = APIRouter(prefix="/api/sessions", tags=["sessions"])
logger = logging.getLogger(__name__)


def parse_version_jsonl(file_path: Path) -> list[dict[str, Any]]:
    """Parse version.jsonl to extract metrics snapshots.

    Each line is a full state snapshot. Extract metrics from the iteration tree.
    Returns list of MetricsSnapshot-compatible dicts.
    """
    snapshots = []

    with open(file_path, "r", encoding="utf-8") as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue

            try:
                state = json.loads(line)
            except json.JSONDecodeError:
                logger.warning(f"Skipping invalid JSON at line {line_num}")
                continue

            iteration = state.get("iteration", {})
            iteration_idx = state.get("iteration_idx", "root")

            # Get current iteration data
            current_iter = iteration.get(iteration_idx, iteration.get("root", {}))

            # Find best candidate with scores
            candidates = current_iter.get("candidates", {})
            best_candidate = None
            best_score = -1

            for cid, cdata in candidates.items():
                score = cdata.get("overall_score")
                if score is not None and score > best_score:
                    best_score = score
                    best_candidate = cdata

            # Also check iteration-level scores (from baseline)
            iter_score = current_iter.get("overall_score")
            if iter_score is not None and (best_candidate is None or iter_score > best_score):
                best_score = iter_score
                best_candidate = {
                    "overall_score": iter_score,
                    "dimension_metrics": {
                        "completeness": current_iter.get("completeness", 0),
                        "clarity_readability": current_iter.get("clarity_readability", 0),
                        "engagement": current_iter.get("engagement", 0),
                        "factual_accuracy": current_iter.get("factual_accuracy", 0),
                        "groundedness": current_iter.get("groundedness", 0),
                    },
                    "experiment_id": current_iter.get("experiment_id", ""),
                    "experiment_link": current_iter.get("experiment_link", ""),
                }

            if best_candidate and best_candidate.get("overall_score") is not None:
                snapshots.append({
                    "version": len(snapshots) + 1,
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                    "overall_score": best_candidate.get("overall_score", 0),
                    "dimensions": best_candidate.get("dimension_metrics", {}),
                    "experiment_id": best_candidate.get("experiment_id", ""),
                    "experiment_link": best_candidate.get("experiment_link", ""),
                })

    return snapshots


def parse_logs_txt(file_path: Path) -> list[dict[str, Any]]:
    """Parse logs.txt to extract agent events.

    Looks for patterns like:
    - "---------- TextMessage (AgentName) ----------"
    - "---------- ToolCallRequestEvent (AgentName) ----------"

    Returns list of AgentEvent-compatible dicts.
    """
    events = []

    # Patterns for message headers
    text_msg_pattern = re.compile(r'^-+ TextMessage \((\w+)\) -+$')
    tool_call_pattern = re.compile(r'^-+ ToolCallRequestEvent \((\w+)\) -+$')
    tool_summary_pattern = re.compile(r'^-+ ToolCallSummaryMessage \((\w+)\) -+$')

    current_agent = None
    current_type = None
    current_content = []

    with open(file_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.rstrip('\n')

            # Check for message headers
            text_match = text_msg_pattern.match(line)
            tool_match = tool_call_pattern.match(line)
            summary_match = tool_summary_pattern.match(line)

            if text_match or tool_match or summary_match:
                # Save previous message if exists
                if current_agent and current_content:
                    content = '\n'.join(current_content).strip()
                    if content:
                        events.append({
                            "type": current_type,
                            "agent": current_agent,
                            "timestamp": datetime.now(timezone.utc).isoformat(),
                            "payload": {"message": content[:2000]},  # Truncate long messages
                        })

                # Start new message
                if text_match:
                    current_agent = text_match.group(1)
                    current_type = "agent_message"
                elif tool_match:
                    current_agent = tool_match.group(1)
                    current_type = "tool_call"
                else:
                    current_agent = summary_match.group(1)
                    current_type = "tool_result"

                current_content = []
            else:
                # Skip noisy status lines
                if line.strip().startswith("Status:"):
                    continue
                current_content.append(line)

        # Don't forget last message
        if current_agent and current_content:
            content = '\n'.join(current_content).strip()
            if content:
                events.append({
                    "type": current_type,
                    "agent": current_agent,
                    "timestamp": datetime.now(timezone.utc).isoformat(),
                    "payload": {"message": content[:2000]},
                })

    return events


def parse_bad_cases(folder_path: Path) -> list[dict[str, Any]]:
    """Parse bad cases from experiment folders.

    Looks for experiments_*/event_details_filtered.jsonl files.
    Returns list of BadCase-compatible dicts.
    """
    bad_cases = []
    seen_ids = set()

    for exp_dir in folder_path.glob("experiments_*"):
        if not exp_dir.is_dir():
            continue

        filtered_file = exp_dir / "event_details_filtered.jsonl"
        if not filtered_file.exists():
            continue

        experiment_id = exp_dir.name.replace("experiments_", "")

        with open(filtered_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue

                try:
                    case = json.loads(line)
                except json.JSONDecodeError:
                    continue

                root_span_id = case.get("root_span_id", "")
                if root_span_id in seen_ids:
                    continue
                seen_ids.add(root_span_id)

                bad_cases.append({
                    "root_span_id": root_span_id,
                    "experiment_id": experiment_id,
                    "input": case.get("test_model_response", "")[:1000],
                    "response": case.get("judge_response", "")[:1000],
                    "reasoning": case.get("judge_reasoning", "")[:1000],
                    "dimension_scores": case.get("merged_scores", {}),
                    "average_score": case.get("average_score", 0),
                    "low_dimensions": case.get("low_dimensions", []),
                })

    return bad_cases


def _get_user_sessions_file(user_id: str) -> Path:
    """Get path to user's sessions.json file."""
    return user_service.get_user_sessions_file(user_id)


def _load_user_sessions(user_id: str) -> list[dict[str, Any]]:
    """Load sessions for a user."""
    sessions_file = _get_user_sessions_file(user_id)
    if not sessions_file.exists():
        return []
    with open(sessions_file, "r", encoding="utf-8") as f:
        return json.load(f)


def _save_user_sessions(user_id: str, sessions: list[dict[str, Any]]) -> None:
    """Save sessions for a user."""
    sessions_file = _get_user_sessions_file(user_id)
    sessions_file.parent.mkdir(parents=True, exist_ok=True)
    with open(sessions_file, "w", encoding="utf-8") as f:
        json.dump(sessions, f, indent=2, default=str)


def _get_session_dir(user_id: str, session_id: str) -> Path:
    """Get path to session data directory."""
    return user_service.get_user_session_dir(user_id, session_id)


def _save_events(session_id: str, events: list[dict[str, Any]]) -> None:
    """Save events to session events file."""
    events_dir = get_store_dir() / "events"
    events_dir.mkdir(parents=True, exist_ok=True)
    events_file = events_dir / f"{session_id}.jsonl"
    with open(events_file, "w", encoding="utf-8") as f:
        for event in events:
            f.write(json.dumps(event, default=str) + "\n")


def _save_versions(session_dir: Path, snapshots: list[dict[str, Any]]) -> None:
    """Save version snapshots to session directory."""
    session_dir.mkdir(parents=True, exist_ok=True)
    versions_file = session_dir / "version.jsonl"
    with open(versions_file, "w", encoding="utf-8") as f:
        for snapshot in snapshots:
            f.write(json.dumps(snapshot, default=str) + "\n")


def _validate_import_folder(folder_path: Path) -> dict[str, bool]:
    """Validate that folder contains required files."""
    return {
        "version.jsonl": (folder_path / "version.jsonl").exists(),
        "logs.txt": (folder_path / "logs.txt").exists(),
    }


@router.post("/import")
async def import_session(
    request: ImportSessionRequest,
    x_user_id: str = Header(..., alias="X-User-Id"),
) -> Session:
    """Import a CLI run folder as a new session."""
    folder_path = Path(request.folder_path)

    if not folder_path.exists():
        raise HTTPException(status_code=404, detail=f"Folder not found: {folder_path}")

    if not folder_path.is_dir():
        raise HTTPException(status_code=400, detail=f"Path is not a directory: {folder_path}")

    # Validate required files
    validation = _validate_import_folder(folder_path)
    missing = [f for f, exists in validation.items() if not exists]
    if missing:
        raise HTTPException(
            status_code=400,
            detail=f"Missing required files: {', '.join(missing)}"
        )

    # Parse data
    logger.info(f"Parsing version.jsonl from {folder_path}")
    snapshots = parse_version_jsonl(folder_path / "version.jsonl")

    logger.info(f"Parsing logs.txt from {folder_path}")
    events = parse_logs_txt(folder_path / "logs.txt")

    logger.info(f"Parsing bad cases from {folder_path}")
    bad_cases = parse_bad_cases(folder_path)

    # Extract metadata from first version snapshot
    task_type = "picasso"
    max_rounds = 20

    # Create session
    session_id = str(int(time.time())) + "_" + str(uuid.uuid4())
    session_name = request.name or f"Imported: {folder_path.name[:20]}"

    session = Session(
        id=session_id,
        user_id=x_user_id,
        name=session_name,
        status=SessionStatus.COMPLETED,
        task_type=task_type,
        created_at=datetime.now(timezone.utc),
        updated_at=datetime.now(timezone.utc),
        current_round=len(snapshots),
        max_rounds=max_rounds,
    )

    # Save session to user's sessions.json
    sessions = _load_user_sessions(x_user_id)
    sessions.append(session.model_dump())
    _save_user_sessions(x_user_id, sessions)

    # Save parsed data
    session_dir = _get_session_dir(x_user_id, session_id)
    _save_events(session_id, events)
    _save_versions(session_dir, snapshots)

    logger.info(f"Imported session {session_id} with {len(snapshots)} versions, {len(events)} events, {len(bad_cases)} bad cases")

    return session
