"""
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
import glob
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 (preserves existing indexing functionality)
        """
        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)

        # Define the standard components of our post-commit hook
        indexing_marker = "Running post-commit hook: Indexing changed files"
        transition_marker = "Checking for test transitions"

        standard_indexing_code = """# Index changed files for RAG
echo "Running post-commit hook: Indexing changed files..."

# Ensure we are in the project root
PROJECT_ROOT=$(git rev-parse --show-toplevel)
cd "$PROJECT_ROOT" || exit 1

# Get list of changed/added Python files in the last commit
# Use --diff-filter=AM to only get Added or Modified files
FILES=$(git diff-tree --no-commit-id --name-only -r HEAD --diff-filter=AM -- "*.py" "*.md" "*.txt")

if [ -z "$FILES" ]; then
  echo "No relevant files changed in this commit."
else
  echo "Files to index:"
  echo "$FILES"

  # Run the indexer using hatch or directly
  if command -v hatch &> /dev/null; then
    # Convert FILES to an argument list
    FILES_ARGS=$(echo "$FILES" | tr '\\n' ' ')
    hatch run python -m chroma_mcp_client.cli index $FILES_ARGS
  else
    # Fallback to direct CLI if hatch not available
    python -m chroma_mcp_client.cli index $FILES
  fi
fi"""

        test_transition_check = """# Added by TestWorkflowManager for test transition tracking
echo "Checking for test transitions..."
python -m chroma_mcp_client.cli check-test-transitions"""

        # Handle post-commit hook (preserve existing content)
        post_commit_path = hooks_dir / "post-commit"

        if post_commit_path.exists():
            # Read existing content
            with open(post_commit_path, "r") as f:
                existing_content = f.read()

            # Check what components we need to add
            indexing_present = indexing_marker in existing_content
            transition_present = transition_marker in existing_content

            # If both are present, no need to modify
            if indexing_present and transition_present:
                logger.info("Both indexing and test transition check already present in post-commit hook")
                os.chmod(post_commit_path, 0o755)  # Still ensure it's executable
                return True

            # Start building new content with the shebang
            if existing_content.startswith("#!/"):
                shebang_line = existing_content.split("\n")[0]
                new_content = f"{shebang_line}\n"
            else:
                new_content = "#!/bin/bash\n"

            # Add auto-generated comment if not present
            if "# Auto-generated by TestWorkflowManager" not in existing_content:
                new_content += "# Auto-generated by TestWorkflowManager\n\n"
            else:
                # Extract and keep the comment line
                for line in existing_content.split("\n"):
                    if "# Auto-generated by TestWorkflowManager" in line:
                        if shebang_line not in new_content:  # Avoid duplicating if we already added it
                            new_content += line + "\n\n"
                        break

            # Add indexing code if needed
            if not indexing_present:
                logger.info("Adding indexing code to post-commit hook")
                new_content += standard_indexing_code + "\n\n"
            else:
                # Extract and preserve existing indexing code
                in_indexing_block = False
                for line in existing_content.split("\n"):
                    if indexing_marker in line:
                        in_indexing_block = True
                        new_content += "# Index changed files for RAG\n"  # Start marker

                    if in_indexing_block:
                        new_content += line + "\n"
                        if "fi" in line and line.strip() == "fi":  # End of the indexing block
                            in_indexing_block = False
                            new_content += "\n"  # Add spacing after block

            # Add test transition check if needed
            if not transition_present:
                logger.info("Adding test transition check to post-commit hook")
                new_content += test_transition_check + "\n"
            else:
                # Extract and preserve existing transition check
                in_transition_block = False
                for line in existing_content.split("\n"):
                    if transition_marker in line:
                        in_transition_block = True
                        if "# Added by TestWorkflowManager for test transition tracking" not in new_content:
                            new_content += "# Added by TestWorkflowManager for test transition tracking\n"

                    if in_transition_block:
                        if (
                            "# Added by TestWorkflowManager for test transition tracking" not in line
                        ):  # Avoid duplication
                            new_content += line + "\n"
                        if "python -m chroma_mcp_client.cli check-test-transitions" in line:
                            in_transition_block = False

            # Extract custom content (anything not part of our standard components)
            custom_content = []
            lines = existing_content.split("\n")
            i = 0
            while i < len(lines):
                line = lines[i]
                skip_line = False

                # Skip shebang and auto-generated lines
                if line.startswith("#!/") or "Auto-generated by TestWorkflowManager" in line:
                    skip_line = True

                # Skip indexing block
                elif indexing_marker in line:
                    skip_line = True
                    # Skip until end of indexing block
                    while i < len(lines) and "fi" not in lines[i]:
                        i += 1
                    if i < len(lines):  # Skip the 'fi' line
                        i += 1
                    continue  # Continue outer loop with updated i

                # Skip transition check lines
                elif transition_marker in line or "# Added by TestWorkflowManager for test transition tracking" in line:
                    skip_line = True
                    # Skip one more line (python command)
                    if i + 1 < len(lines) and "python -m chroma_mcp_client.cli check-test-transitions" in lines[i + 1]:
                        i += 1

                if not skip_line:
                    custom_content.append(line)
                i += 1

            # Add custom content if any was found
            if custom_content and any(line.strip() for line in custom_content):
                new_content += "\n# Custom hook content (preserved)\n"
                new_content += "\n".join(custom_content).strip() + "\n"

            # Write the updated content
            with open(post_commit_path, "w") as f:
                f.write(new_content)
            logger.info("Updated post-commit hook")
        else:
            # Create new post-commit hook with all standard components
            standard_post_commit = f"""#!/bin/bash
# Auto-generated by TestWorkflowManager

{standard_indexing_code}

{test_transition_check}
"""
            with open(post_commit_path, "w") as f:
                f.write(standard_post_commit)
            logger.info("Created new post-commit hook with indexing and test transition check")

        # Ensure it's executable
        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 cleanup_processed_artifacts(self, workflow_file: str) -> bool:
        """
        Clean up test artifacts after successful processing.

        Args:
            workflow_file: Path to the workflow file that was processed

        Returns:
            bool: True if cleanup was successful, False otherwise
        """
        try:
            logger.info(f"Cleaning up processed artifacts from workflow: {workflow_file}")

            # Read workflow file to get associated artifact paths
            if not os.path.exists(workflow_file):
                logger.warning(f"Workflow file not found: {workflow_file}")
                return False

            with open(workflow_file, "r") as f:
                workflow_data = json.load(f)

            # Get paths to artifacts
            before_xml = workflow_data.get("before_xml")
            after_xml = workflow_data.get("after_xml")
            status = workflow_data.get("status")

            if status != "transitioned":
                logger.warning(
                    f"Cannot clean up workflow with status '{status}'. Only 'transitioned' workflows can be cleaned up."
                )
                return False

            files_removed = 0

            # Only remove files if they exist
            if before_xml and os.path.exists(before_xml):
                logger.info(f"Removing processed artifact: {before_xml}")
                os.remove(before_xml)
                files_removed += 1

                # Also remove commit file if it exists
                commit_file = f"{before_xml}.commit"
                if os.path.exists(commit_file):
                    os.remove(commit_file)
                    files_removed += 1

            # We can optionally remove the "after" XML as well, since it's already logged
            # Uncomment if desired (may be useful to keep latest test results)
            # if after_xml and os.path.exists(after_xml):
            #     logger.info(f"Removing processed artifact: {after_xml}")
            #     os.remove(after_xml)
            #     files_removed += 1

            # Only remove the workflow file itself after processing
            logger.info(f"Removing processed workflow file: {workflow_file}")
            os.remove(workflow_file)
            files_removed += 1

            # Find and remove any old workflow file that points to the same before_xml
            workflow_dir = os.path.dirname(workflow_file)
            old_workflow_pattern = os.path.join(workflow_dir, "test_workflow_*.json")
            for old_workflow_file in glob.glob(old_workflow_pattern):
                try:
                    with open(old_workflow_file, "r") as f:
                        old_data = json.load(f)
                    if old_data.get("xml_path") == before_xml:
                        logger.info(f"Removing related workflow file: {old_workflow_file}")
                        os.remove(old_workflow_file)
                        files_removed += 1
                except Exception as e:
                    logger.warning(f"Error processing old workflow file {old_workflow_file}: {e}")

            logger.info(f"Cleanup completed successfully. Removed {files_removed} files.")
            return True

        except Exception as e:
            logger.warning(f"Error cleaning up artifacts: {e}")
            return False


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()


def cleanup_test_artifacts(workflow_file: str) -> bool:
    """
    Clean up test artifacts from a completed workflow.

    Args:
        workflow_file: Path to the workflow file that tracks the artifacts

    Returns:
        bool: True if cleanup was successful, False otherwise
    """
    try:
        workspace_dir = os.path.abspath(os.path.join(os.path.dirname(workflow_file), "../.."))
        workflow_manager = TestWorkflowManager(workspace_dir=workspace_dir)
        return workflow_manager.cleanup_processed_artifacts(workflow_file)
    except Exception as e:
        logger.error(f"Error cleaning up test artifacts: {e}")
        return False
