#!/usr/bin/env python3
"""
Primer Generation using Prompt Dependency Resolution

This script demonstrates the full workflow of generating InfoTech primers
using MarkiTect's Prompt Dependency Resolution infrastructure.

Features demonstrated:
- Artifact creation and storage
- PromptTemplate with macro resolution
- Dependency tracking
- Quality validation
- Incremental recomputation
- Full provenance tracing
"""

import sys
from pathlib import Path
from typing import Optional

# Add project root to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from markitect.prompts.models import Artifact, ArtifactType
from markitect.prompts.repositories.sqlite import SQLiteArtifactRepository
from markitect.prompts.dependencies.repository import SQLiteDependencyRepository
from markitect.prompts.services.artifact_service import ArtifactService
from markitect.prompts.templates.models import PromptTemplate, ContentMacro, MacroKind
from markitect.prompts.templates.analyzer import TemplateAnalyzer
from markitect.prompts.resolver.resolver import PromptResolver
from markitect.prompts.resolver.compiler import ContextCompiler
from markitect.prompts.resolver.strategy import ResolutionConfig, MultiSpaceResolutionStrategy
from markitect.prompts.execution.models import PromptRun, RunConfig
from markitect.prompts.execution.manifest import RunManifest
from markitect.prompts.dependencies.graph import GraphBuilder
from markitect.prompts.traceability.service import TraceabilityService
from markitect.prompts.queries.operations import PromptQueryService
from markitect.prompts.visualization.graph import GraphExporter


class PrimerGenerator:
    """Generates InfoTech primers using prompt dependency resolution."""

    def __init__(self, db_path: str = "primers.db"):
        """Initialize with database path."""
        self.db_path = db_path
        self.artifact_repo = SQLiteArtifactRepository(db_path)
        self.artifact_service = ArtifactService(self.artifact_repo)
        self.dep_repo = SQLiteDependencyRepository(db_path)
        self.graph_builder = GraphBuilder(self.dep_repo)
        self.trace_service = TraceabilityService(
            self.artifact_repo, self.dep_repo, db_path=db_path
        )
        self.query_service = PromptQueryService(
            self.artifact_repo, self.dep_repo, db_path=db_path
        )

        # Create information spaces
        self.spaces = {
            "templates": "primer-templates",
            "topics": "primer-topics",
            "guidelines": "primer-guidelines",
            "output": "generated-primers",
        }

    def load_or_create_artifact(
        self, space: str, filepath: Path, artifact_type: ArtifactType, name: Optional[str] = None
    ) -> tuple[Artifact, str]:
        """Load artifact from file or fetch from repository. Returns (artifact, content)."""
        if name is None:
            name = filepath.stem

        # Load content from file
        if not filepath.exists():
            raise FileNotFoundError(f"Artifact file not found: {filepath}")

        content = filepath.read_text()

        # Check if artifact exists
        existing = self.artifact_repo.get_by_name(space, name)
        if existing:
            print(f"✓ Found existing artifact: {name}")
            return existing, content

        # Create and store artifact
        artifact = Artifact.create(
            space_id=space,
            name=name,
            content=content,
            artifact_type=artifact_type,
        )
        artifact = self.artifact_repo.create(artifact)
        print(f"✓ Created artifact: {name} (digest: {artifact.content_digest[:8]})")
        return artifact, content

    def setup_artifacts(self, example_dir: Path):
        """Load all artifacts into the repository."""
        print("\n=== Loading Artifacts ===")

        # Cache for artifact content (repository doesn't store content)
        self.artifact_content = {}

        # Load template
        template_file = example_dir / "templates" / "generate-primer.md"
        self.template_artifact, content = self.load_or_create_artifact(
            self.spaces["templates"],
            template_file,
            ArtifactType.TEMPLATE,
        )
        self.artifact_content[self.template_artifact.id] = content

        # Load topic artifacts
        topics_dir = example_dir / "artifacts" / "topics"
        self.topic_artifacts = {}
        for topic_file in topics_dir.glob("*.md"):
            artifact, content = self.load_or_create_artifact(
                self.spaces["topics"],
                topic_file,
                ArtifactType.CONTENT,
            )
            self.topic_artifacts[artifact.name] = artifact
            self.artifact_content[artifact.id] = content

        # Load guideline artifacts (rename to match macro expectations)
        guidelines_dir = example_dir / "artifacts" / "guidelines"
        self.guideline_artifacts = {}
        guideline_name_map = {
            "authoring-rules.md": "authoring_rules",
            "research-prompt.md": "research_prompt",
        }
        for guide_file in guidelines_dir.glob("*.md"):
            # Use mapped name if available, otherwise use stem
            artifact_name = guideline_name_map.get(guide_file.name, guide_file.stem)
            artifact, content = self.load_or_create_artifact(
                self.spaces["guidelines"],
                guide_file,
                ArtifactType.CONTENT,
                name=artifact_name,
            )
            self.guideline_artifacts[artifact.name] = artifact
            self.artifact_content[artifact.id] = content

    def create_template(self) -> PromptTemplate:
        """Create PromptTemplate from template artifact."""
        print("\n=== Parsing Template ===")

        # Create PromptTemplate from artifact
        template = PromptTemplate.from_artifact(self.template_artifact)

        # For demo: manually add macros since analyzer would extract from {{...}}
        # In real usage, analyzer would parse the content automatically
        template.macros = [
            ContentMacro(kind=MacroKind.REQUIRED, target="topic"),
            ContentMacro(kind=MacroKind.REQUIRED, target="authoring_rules"),
            ContentMacro(kind=MacroKind.REQUIRED, target="research_prompt"),
        ]
        template.analyzed = True

        print(f"✓ Template created with {len(template.macros)} macro dependencies")
        return template

    def generate_primer(self, topic_name: str, output_name: Optional[str] = None) -> Artifact:
        """Generate a primer for the given topic."""
        print(f"\n=== Generating Primer: {topic_name} ===")

        # Get topic artifact
        topic_artifact = self.topic_artifacts.get(topic_name)
        if not topic_artifact:
            raise ValueError(f"Topic not found: {topic_name}")

        # Get the topic's content from cache
        topic_content = self.artifact_content.get(topic_artifact.id)
        if not topic_content:
            raise ValueError(f"Topic content not found for: {topic_name}")

        # Create a temporary "topic" artifact with this specific topic's content
        # This allows the template's @{topic} macro to resolve
        temp_topic = Artifact.create(
            space_id=self.spaces["topics"],
            name="topic",
            content=topic_content,
            artifact_type=ArtifactType.CONTENT,
        )

        # Check if "topic" artifact already exists and delete it
        existing_topic = self.artifact_repo.get_by_name(self.spaces["topics"], "topic")
        if existing_topic:
            self.artifact_repo.delete(existing_topic.id)

        temp_topic = self.artifact_repo.create(temp_topic)
        self.artifact_content[temp_topic.id] = topic_content
        print(f"✓ Bound topic macro to {topic_name}")

        # Create template
        template = self.create_template()

        # Configure resolution
        resolution_config = ResolutionConfig(
            space_id=self.spaces["templates"],
            included_spaces=[
                self.spaces["topics"],
                self.spaces["guidelines"],
            ],
        )

        print("\n=== Resolving Dependencies ===")
        strategy = MultiSpaceResolutionStrategy()
        resolver = PromptResolver(self.artifact_service, strategy)
        resolution_result = resolver.resolve_template(template, resolution_config)

        if not resolution_result.success:
            print(f"✗ Resolution failed: {resolution_result.context.errors}")
            return None

        print(f"✓ Resolved {len(resolution_result.context.resolved_macros)} macros")
        for resolved in resolution_result.context.resolved_macros:
            print(f"  - {resolved.macro.target} → {resolved.artifact.name}")

        # Compile the prompt
        print("\n=== Compiling Prompt ===")
        compiler = ContextCompiler()

        # For demo, use mock template content with macro markers
        template_content = f"""
# Generate InfoTech Primer

Topic: @{{topic}}
Guidelines: @{{authoring_rules}}
Research Protocol: @{{research_prompt}}

Generate a complete primer following the authoring rules.
"""

        compiled = compiler.compile(template, template_content, resolution_result)
        print(f"✓ Compiled prompt (digest: {compiled.content_digest[:8]})")

        # Create run manifest
        run_id = f"run-{topic_name}-001"
        manifest = RunManifest.create(
            run_id=run_id,
            template_id=template.artifact.id,
            template_name=template.artifact.name,
            template_digest=self.template_artifact.content_digest,
        )

        # Add resolved inputs
        for resolved in resolution_result.context.resolved_macros:
            manifest.add_resolved_input(
                name=resolved.macro.target,
                artifact_id=resolved.artifact.id,
                space_id=resolved.space_id,
                digest=resolved.artifact.content_digest,
            )

            # Create dependency edge
            manifest.add_dependency_edge(
                source_id=resolved.artifact.id,
                target_id=run_id,
                edge_type="requires",
            )

        # For demo: create output artifact (would normally come from LLM)
        output_name = output_name or f"{topic_name}-primer"
        output_content = f"""# {topic_name.upper()} Primer

[Generated primer content would go here...]

This primer was generated using the Prompt Dependency Resolution infrastructure.

**Dependencies:**
- Template: {template.artifact.name}
- Topic: {topic_artifact.name}
- Guidelines: authoring-rules, research-prompt

**Run ID:** {run_id}
"""

        output_artifact = Artifact.create(
            space_id=self.spaces["output"],
            name=output_name,
            content=output_content,
            artifact_type=ArtifactType.GENERATED,
        )
        output_artifact = self.artifact_repo.create(output_artifact)

        # Add output to manifest
        manifest.add_output_artifact(
            artifact_id=output_artifact.id,
            name=output_artifact.name,
            digest=output_artifact.content_digest,
            artifact_type=output_artifact.artifact_type.value,
        )

        manifest.add_dependency_edge(
            source_id=run_id,
            target_id=output_artifact.id,
            edge_type="generates",
        )

        # Persist all dependency edges at once
        edges = self.graph_builder.persist_edges(manifest)
        print(f"✓ Persisted {len(edges)} dependency edges")

        print(f"✓ Generated primer: {output_artifact.name}")
        print(f"  ID: {output_artifact.id}")
        print(f"  Digest: {output_artifact.content_digest[:8]}")

        return output_artifact

    def show_provenance(self, artifact_id: str):
        """Display full provenance trace for an artifact."""
        print(f"\n=== Provenance Trace ===")

        trace = self.trace_service.trace_artifact(artifact_id)

        print(f"Artifact: {trace.artifact_id}")

        if trace.producing_run:
            print(f"\nProducing Run: {trace.producing_run.run_id}")
            print(f"  Template: {trace.producing_run.template_id}")
            print(f"  Status: {trace.producing_run.status}")

        if trace.input_artifacts:
            print(f"\nInput Artifacts ({len(trace.input_artifacts)}):")
            for inp in trace.input_artifacts:
                print(f"  - {inp.name} ({inp.artifact_type})")

        if trace.dependency_chain:
            print(f"\nDependency Chain ({len(trace.dependency_chain)} artifacts):")
            for dep_id in trace.dependency_chain[:5]:  # Show first 5
                print(f"  - {dep_id[:8]}...")

    def visualize_dependencies(self, artifact_id: str, output_file: str = "dependencies.mermaid"):
        """Generate dependency graph visualization."""
        print(f"\n=== Dependency Visualization ===")

        from markitect.prompts.dependencies.queries import DependencyQueryService

        query_svc = DependencyQueryService(self.dep_repo)
        deps = query_svc.find_transitive_dependencies(artifact_id)
        dependents = query_svc.find_transitive_dependents(artifact_id)
        all_ids = deps | dependents | {artifact_id}

        graph = self.graph_builder.build_graph(all_ids)
        mermaid = GraphExporter.to_mermaid(graph, f"Primer Generation Dependencies")

        Path(output_file).write_text(mermaid)
        print(f"✓ Saved dependency graph to {output_file}")


def main():
    """Main execution flow."""
    print("╔══════════════════════════════════════════════════════════════╗")
    print("║  MarkiTect Prompt Dependency Resolution Example             ║")
    print("║  InfoTech Primer Generation                                  ║")
    print("╚══════════════════════════════════════════════════════════════╝")

    example_dir = Path(__file__).parent

    # Initialize generator
    generator = PrimerGenerator(db_path=str(example_dir / "primers.db"))

    # Setup artifacts
    generator.setup_artifacts(example_dir)

    # Generate primers for multiple topics
    topics = ["etl", "microservices"]

    generated_artifacts = []
    for topic in topics:
        try:
            artifact = generator.generate_primer(topic)
            if artifact:
                generated_artifacts.append(artifact)
        except Exception as e:
            print(f"✗ Failed to generate primer for {topic}: {e}")

    # Show provenance for first generated primer
    if generated_artifacts:
        artifact = generated_artifacts[0]
        generator.show_provenance(artifact.id)
        generator.visualize_dependencies(artifact.id)

    # Show statistics
    print("\n=== Dependency Statistics ===")
    stats = generator.query_service.get_dependency_stats()
    print(f"Total nodes: {stats['total_nodes']}")
    print(f"Total edges: {stats['total_edges']}")
    print(f"Root artifacts: {stats['root_count']}")
    print(f"Leaf artifacts: {stats['leaf_count']}")
    print(f"Has cycles: {stats['has_cycles']}")

    print("\n✓ Primer generation complete!")
    print(f"  Database: {generator.db_path}")
    print(f"  Generated: {len(generated_artifacts)} primers")


if __name__ == "__main__":
    main()
