"""
Automated Test-Driven Learning Workflow

This module implements the automated workflow for capturing test transitions,
linking them to code changes and chat history, and promoting validated learning.

The workflow automates these steps:
1. Capture test failures with full context
2. Monitor for subsequent test success after code changes
3. Create validation evidence with bidirectional links
4. Generate derived learnings from validated fixes
"""

import os
import json
import logging
import datetime
import subprocess
import tempfile
import uuid
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any, Set

from .test_collector import parse_junit_xml, create_test_transition_evidence, store_test_results
from .schemas import TestTransitionEvidence, ValidationEvidence
from .evidence_collector import collect_validation_evidence
from ..learnings import promote_to_learnings_collection
from chroma_mcp.utils.chroma_client import get_chroma_client

logger = logging.getLogger(__name__)


class TestWorkflowManager:
    """
    Manages the automated test-driven learning workflow.

    This class handles the complete lifecycle of test-driven learning:
    - Initial test failure capture
    - Monitoring for fixes
    - Correlation with chat history
    - Auto-promotion of validated learnings
    """

    def __init__(
        self,
        workspace_dir: str = ".",
        test_results_collection: str = "test_results_v1",
        chat_history_collection: str = "chat_history_v1",
        evidence_collection: str = "validation_evidence_v1",
        temp_dir: Optional[str] = None,
        chroma_client=None,
    ):
        """
        Initialize the test workflow manager.

        Args:
            workspace_dir: Root directory of the workspace
            test_results_collection: ChromaDB collection for test results
            chat_history_collection: ChromaDB collection for chat history
            evidence_collection: ChromaDB collection for validation evidence
            temp_dir: Directory for temporary files (defaults to system temp)
            chroma_client: Optional ChromaDB client instance
        """
        self.workspace_dir = Path(workspace_dir).absolute()
        self.test_results_collection = test_results_collection
        self.chat_history_collection = chat_history_collection
        self.evidence_collection = evidence_collection
        self.temp_dir = temp_dir or tempfile.gettempdir()

        # Get ChromaDB client
        self.client = chroma_client or get_chroma_client()

        # Ensure collections exist
        self._ensure_collections()

    def _ensure_collections(self):
        """Ensure required ChromaDB collections exist."""
        for collection_name in [self.test_results_collection, self.chat_history_collection, self.evidence_collection]:
            try:
                self.client.get_collection(name=collection_name)
                logger.info(f"Collection {collection_name} exists")
            except Exception:
                self.client.create_collection(name=collection_name)
                logger.info(f"Created collection {collection_name}")

    def setup_git_hooks(self):
        """
        Set up git hooks to track test results.

        This creates:
        - pre-push hook to run tests and capture results
        - post-commit hook to check for test transitions
        """
        hooks_dir = self.workspace_dir / ".git" / "hooks"

        if not hooks_dir.exists():
            logger.warning(f"Git hooks directory not found: {hooks_dir}")
            return False

        # Create pre-push hook
        pre_push_path = hooks_dir / "pre-push"
        pre_push_content = """#!/bin/bash
# Auto-generated by TestWorkflowManager
echo "Running tests before push..."
./scripts/test.sh -c -v --auto-capture-workflow
"""
        with open(pre_push_path, "w") as f:
            f.write(pre_push_content)
        os.chmod(pre_push_path, 0o755)

        # Create post-commit hook
        post_commit_path = hooks_dir / "post-commit"
        post_commit_content = """#!/bin/bash
# Auto-generated by TestWorkflowManager
echo "Checking for test transitions..."
python -m chroma_mcp_client.cli check-test-transitions
"""
        with open(post_commit_path, "w") as f:
            f.write(post_commit_content)
        os.chmod(post_commit_path, 0o755)

        logger.info("Git hooks set up successfully")
        return True

    def capture_test_failure(self, xml_path: str, commit_hash: Optional[str] = None) -> str:
        """
        Capture a test failure and store it for future comparison.

        Args:
            xml_path: Path to JUnit XML with test results
            commit_hash: Optional git commit hash

        Returns:
            ID of the stored test run
        """
        try:
            # Parse and store test results
            results = parse_junit_xml(xml_path)

            # Get current git commit hash if not provided
            if not commit_hash:
                try:
                    commit_hash = (
                        subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=self.workspace_dir).decode().strip()
                    )
                except Exception as e:
                    logger.warning(f"Failed to get git commit hash: {e}")
                    commit_hash = "unknown"

            # Store test results in ChromaDB
            run_id = store_test_results(
                results_dict=results, collection_name=self.test_results_collection, chroma_client=self.client
            )

            # Save metadata about this test run for future reference
            run_metadata = {
                "run_id": run_id,
                "timestamp": datetime.datetime.now().isoformat(),
                "commit_hash": commit_hash,
                "failing_tests": [
                    test_id for test_id, result in results.items() if result["status"] in ("fail", "error")
                ],
                "workflow_stage": "initial_failure",
            }

            # Save metadata to a JSON file
            metadata_path = Path(self.temp_dir) / f"test_run_{run_id}.json"
            with open(metadata_path, "w") as f:
                json.dump(run_metadata, f)

            logger.info(f"Captured test failure with run ID: {run_id}")
            return run_id

        except Exception as e:
            logger.error(f"Failed to capture test failure: {e}")
            raise

    def find_chat_sessions_for_code_changes(self, before_commit: str, after_commit: str) -> List[str]:
        """
        Find chat history sessions that contributed to code changes.

        Args:
            before_commit: Git commit hash from before changes
            after_commit: Git commit hash from after changes

        Returns:
            List of chat history session IDs that modified relevant code
        """
        # Get collection
        try:
            chat_collection = self.client.get_collection(name=self.chat_history_collection)
        except Exception as e:
            logger.error(f"Failed to get chat history collection: {e}")
            return []

        # Get list of changed files between commits
        try:
            changed_files = (
                subprocess.check_output(
                    ["git", "diff", "--name-only", before_commit, after_commit], cwd=self.workspace_dir
                )
                .decode()
                .strip()
                .split("\n")
            )
        except Exception as e:
            logger.warning(f"Failed to get changed files: {e}")
            return []

        if not changed_files or changed_files[0] == "":
            logger.warning("No files changed between commits")
            return []

        # Query chat history for entries that modified these files
        chat_ids = set()
        for file_path in changed_files:
            try:
                results = chat_collection.query(
                    query_texts=[f"modified {file_path}"],
                    n_results=10,
                    where={
                        "$or": [
                            {"tool_sequence": {"$contains": "edit_file"}},
                            {"tool_sequence": {"$contains": "search_replace"}},
                        ]
                    },
                )

                if results and results.get("ids"):
                    chat_ids.update(results["ids"][0])
            except Exception as e:
                logger.warning(f"Failed to query chat history for {file_path}: {e}")
                continue

        return list(chat_ids)

    def create_validation_from_test_transition(
        self, before_xml: str, after_xml: str, before_commit: Optional[str] = None, after_commit: Optional[str] = None
    ) -> Tuple[ValidationEvidence, List[str]]:
        """
        Create validation evidence from a test transition.

        Args:
            before_xml: Path to JUnit XML from before changes
            after_xml: Path to JUnit XML from after changes
            before_commit: Optional git commit hash from before
            after_commit: Optional git commit hash from after

        Returns:
            Tuple of validation evidence and list of related chat IDs
        """
        # Create test transition evidence
        transitions = create_test_transition_evidence(
            before_xml=before_xml, after_xml=after_xml, commit_before=before_commit, commit_after=after_commit
        )

        # If we have commit hashes, find related chat sessions
        chat_ids = []
        if before_commit and after_commit:
            chat_ids = self.find_chat_sessions_for_code_changes(before_commit=before_commit, after_commit=after_commit)

        # Create validation evidence
        evidence = collect_validation_evidence(
            test_transitions=transitions, runtime_errors=[], code_quality_improvements=[]
        )

        # Add related chat IDs to evidence
        for transition in evidence.test_transitions:
            transition.related_chat_ids = chat_ids

        # Store evidence in ChromaDB if score is significant
        if evidence.meets_threshold():
            # TODO: Implement storage logic
            pass

        return evidence, chat_ids

    def auto_promote_learning(
        self, evidence: ValidationEvidence, chat_ids: List[str], confidence_threshold: float = 0.8
    ) -> Optional[str]:
        """
        Automatically promote a learning based on validation evidence.

        Args:
            evidence: Validation evidence
            chat_ids: List of related chat history IDs
            confidence_threshold: Threshold for auto-promotion

        Returns:
            ID of promoted learning if successful, None otherwise
        """
        if not evidence.meets_threshold() or evidence.score < confidence_threshold:
            logger.info(f"Evidence score {evidence.score} below threshold {confidence_threshold}")
            return None

        if not chat_ids:
            logger.warning("No chat IDs available for learning promotion")
            return None

        # TODO: Implement auto-promotion logic
        # This would extract the key information from the chat history
        # and create a new entry in derived_learnings_v1

        return None


def check_for_completed_workflows():
    """
    Check for completed test-driven learning workflows.

    This function looks for:
    1. Previously captured test failures
    2. New passing test results for the same tests
    3. Creates validation evidence and potentially promotes learnings

    Returns:
        Number of workflows processed
    """
    # TODO: Implement workflow detection and processing
    return 0


def setup_automated_workflow(workspace_dir: str = "."):
    """
    Set up the automated test-driven learning workflow.

    Args:
        workspace_dir: Root directory of the workspace

    Returns:
        True if setup succeeded, False otherwise
    """
    manager = TestWorkflowManager(workspace_dir=workspace_dir)
    return manager.setup_git_hooks()
