"""Pipeline runners.

High-level functions that compose pipeline stages into complete workflows:
process_image_unified, process_tex_item, process_generated_problem.
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, Optional

from vbagent.cli.common import _get_console, _get_panel, format_latex, extract_problem_solution
from vbagent.pipeline.io import (
    merge_metadata_into_latex,
    insert_tikz_into_latex,
    convert_primary_to_classification,
)
from vbagent.pipeline.stages import (
    generate_solution_orchestrated,
    generate_problem_diagrams,
    assess_difficulty_stage,
    extract_ideas_stage,
    generate_alternate_stage,
    generate_variants_stage,
)

if TYPE_CHECKING:
    from vbagent.models.pipeline import PipelineResult


def process_tex_item(
    tex_content: str,
    source_path: str,
    variant_types: list[str],
    generate_alternate: bool,
    generate_ideas: bool = False,
    use_context: bool = True,
) -> "PipelineResult":
    """Process a TeX item through the pipeline (skips classification/scanning)."""
    from vbagent.agents.content_generation.idea import extract_ideas
    from vbagent.agents.variants.variant import generate_variant
    from vbagent.models.classification import ClassificationResult
    from vbagent.models.pipeline import PipelineResult
    from vbagent.config import get_config

    console = _get_console()

    classification = ClassificationResult(
        subject=get_config().subject,
        question_type="subjective",
        has_diagram=False,
        confidence=1.0,
        classified_from="latex",
    )

    latex = tex_content
    problem, solution = extract_problem_solution(latex)

    # Ideas
    ideas = None
    if generate_ideas and problem and solution:
        with console.status("[bold green]Extracting ideas..."):
            ideas = extract_ideas(problem, solution)
        ideas_text = f"[bold]Concepts:[/bold] {', '.join(ideas.concepts)}\n"
        ideas_text += f"[bold]Formulas:[/bold] {', '.join(ideas.formulas)}\n"
        ideas_text += f"[bold]Techniques:[/bold] {', '.join(ideas.techniques)}\n"
        ideas_text += f"[bold]Difficulty Factors:[/bold] {', '.join(ideas.difficulty_factors)}"
        console.print(_get_panel(ideas_text, title="Extracted Ideas", border_style="yellow"))

    # Alternates
    alternate_solutions = []
    if generate_alternate and problem and solution:
        with console.status("[bold green]Generating alternate solution..."):
            alt = generate_alternate_solution(problem, solution, ideas)
            alternate_solutions.append(alt)
        console.print(_get_panel(alt, title="Alternate Solution", border_style="magenta"))

    # Variants
    variants = {}
    for vtype in variant_types:
        with console.status(f"[bold green]Generating {vtype} variant..."):
            variant_latex = generate_variant(latex, vtype, ideas, use_context=use_context)
            variants[vtype] = variant_latex
        console.print(_get_panel(variant_latex, title=f"{vtype.title()} Variant", border_style="green"))

    return PipelineResult(
        source_path=source_path,
        classification=classification,
        latex=latex,
        tikz_code=None,
        ideas=ideas,
        alternate_solutions=alternate_solutions,
        variants=variants,
    )


def generate_alternate_solution(problem: str, solution: str, ideas) -> str:
    """Generate an alternate solution using the alternate agent."""
    from vbagent.agents.content_generation.alternate import generate_alternate
    return generate_alternate(problem, solution, ideas)


def process_generated_problem(
    generated,
    problem_num: int,
    output_base_dir: Path = Path("agentic"),
) -> dict:
    """Process a generated problem through the full pipeline.

    Runs the complete workflow on a problem generated by Agent 5:
    1. Generate TikZ from diagram_description (if present)
    2. Run Agent 4 (LaTeX Classifier) on the problem
    3. Run Agent 2 (Diagram Analyzer) if has_diagram
    4. Run Agent 3 (Difficulty Assessor)
    5. Merge metadata into LaTeX
    6. Save everything in standard format
    """
    from vbagent.agents.classification.latex_classifier import classify_from_latex
    from vbagent.agents.classification.diagram_analyzer import analyze_diagram_from_description
    from vbagent.agents.classification.difficulty_assessor import assess_difficulty
    from vbagent.agents.diagram.tikz import generate_tikz
    import json

    console = _get_console()

    dirs = {
        "generated": output_base_dir / "generated",
        "classifications": output_base_dir / "classifications",
        "diagrams": output_base_dir / "diagrams",
        "difficulty": output_base_dir / "difficulty",
        "tikz": output_base_dir / "tikz",
    }
    for dir_path in dirs.values():
        dir_path.mkdir(parents=True, exist_ok=True)

    problem_name = f"problem_{problem_num}"

    # Step 1: Generate TikZ if diagram_description exists
    tikz_code = None
    if generated.diagram_description:
        console.print("[cyan]→ Generating TikZ from description...[/cyan]")
        try:
            tikz_code = generate_tikz(description=generated.diagram_description, use_context=True)
            tikz_file = dirs["tikz"] / f"{problem_name}.tex"
            tikz_file.write_text(tikz_code)
            console.print(f"[green]✓ TikZ saved to {tikz_file}[/green]")
            if r'\input{diagram}' in generated.problem_latex:
                generated.problem_latex = insert_tikz_into_latex(generated.problem_latex, tikz_code)
        except Exception as e:
            console.print(f"[yellow]! TikZ generation failed: {e}[/yellow]")

    # Step 2: Classify
    console.print("[cyan]→ Classifying generated problem...[/cyan]")
    primary = classify_from_latex(generated.problem_latex)
    classification_file = dirs["classifications"] / f"{problem_name}.json"
    classification_file.write_text(json.dumps(primary.model_dump(), indent=2))
    console.print("[green]✓ Classification saved[/green]")

    # Step 3: Diagram analysis
    diagram = None
    if primary.has_diagram or generated.diagram_description:
        console.print("[cyan]→ Analyzing diagram...[/cyan]")
        try:
            diagram = analyze_diagram_from_description(
                description=generated.diagram_description or "Generated diagram",
                primary=primary,
            )
            diagram_file = dirs["diagrams"] / f"{problem_name}.json"
            diagram_file.write_text(json.dumps(diagram.model_dump(), indent=2))
            console.print("[green]✓ Diagram analysis saved[/green]")
        except Exception as e:
            console.print(f"[yellow]! Diagram analysis failed: {e}[/yellow]")

    # Step 4: Difficulty assessment
    console.print("[cyan]→ Assessing difficulty...[/cyan]")
    difficulty = assess_difficulty(
        latex_content=generated.problem_latex + "\n\n" + generated.solution_latex,
        primary=primary,
        diagram=diagram,
    )
    difficulty_file = dirs["difficulty"] / f"{problem_name}.json"
    difficulty_file.write_text(json.dumps(difficulty.model_dump(), indent=2))
    console.print("[green]✓ Difficulty assessment saved[/green]")

    # Step 5: Merge metadata
    console.print("[cyan]→ Merging metadata...[/cyan]")
    latex_with_metadata = merge_metadata_into_latex(
        latex=generated.problem_latex, primary=primary, diagram=diagram, difficulty=difficulty,
    )

    # Step 6: Save
    problem_file = dirs["generated"] / f"{problem_name}.tex"
    problem_file.write_text(latex_with_metadata)
    console.print(f"[green]✓ Problem saved to {problem_file}[/green]")

    solution_file = dirs["generated"] / f"{problem_name}_solution.tex"
    solution_file.write_text(generated.solution_latex)

    idea_file = dirs["generated"] / f"{problem_name}_idea.tex"
    idea_file.write_text(generated.idea_latex)

    return {
        "problem_path": str(problem_file),
        "solution_path": str(solution_file),
        "idea_path": str(idea_file),
        "tikz_path": str(dirs["tikz"] / f"{problem_name}.tex") if tikz_code else None,
        "classification": primary.model_dump(),
        "diagram_analysis": diagram.model_dump() if diagram else None,
        "difficulty": difficulty.model_dump(),
    }


def process_image_unified(
    image_path: str,
    variant_types: list[str] = None,
    generate_alternate: bool = False,
    generate_ideas: bool = False,
    use_context: bool = True,
    assess_difficulty: bool = False,
    merge_metadata: bool = True,
    use_cache: bool = True,
    problem_id: Optional[str] = None,
    use_orchestrator: bool = False,
    generate_solution: bool = False,
) -> "PipelineResult":
    """Process an image using the unified pipeline (fewer API calls).

    Architecture:
    1. Unified Classifier (1 call → subject + type + diagram analysis)
    2. Problem Orchestrator (scan ∥ tikz, deterministic routing)
    3. Optional: Solution Orchestrator, Ideas, Alternates, Variants
    """
    from vbagent.cache import PipelineCache
    from vbagent.models.pipeline import PipelineResult
    from vbagent.models.classification import ClassificationResult
    from vbagent.pipeline.stages import (
        classify_unified,
        run_problem_orchestrator,
    )
    from vbagent.agents.classification.unified_classifier import to_primary, to_diagram_analysis

    variant_types = variant_types or []
    console = _get_console()

    cache = PipelineCache() if use_cache else None
    if problem_id is None:
        problem_id = Path(image_path).stem

    if cache:
        cached_stages = cache.get_cached_stages(problem_id)
        if cached_stages:
            console.print(f"[dim]Cached: {', '.join(cached_stages)}[/dim]")

    # Stage 1: Unified classification (1 API call)
    classification = classify_unified(
        image_path, cache=cache, problem_id=problem_id, console=console,
    )
    primary = to_primary(classification)
    diagram_analysis = to_diagram_analysis(classification)

    console.print(f"[cyan]Subject:[/cyan] {classification.subject}")
    console.print(f"[cyan]Type:[/cyan] {classification.question_type}")
    console.print(f"[cyan]Diagram:[/cyan] {'Yes' if classification.has_diagram else 'No'}"
                  + (f" ({classification.diagram_type})" if classification.diagram_type else ""))

    # Stage 2: Problem Orchestrator (scan + tikz, parallel)
    problem_result = run_problem_orchestrator(
        image_path, classification,
        use_context=use_context, cache=cache, problem_id=problem_id, console=console,
    )
    latex = problem_result.latex
    tikz_code = problem_result.tikz_code

    # Stage 2b: Solution Orchestrator (optional)
    if use_orchestrator or generate_solution:
        latex = generate_solution_orchestrated(
            image_path, primary, problem_latex=latex,
            cache=cache, problem_id=problem_id, console=console,
        )

    # Stage 3: Difficulty assessment (optional)
    difficulty_result = None
    if assess_difficulty:
        difficulty_result = assess_difficulty_stage(
            latex, primary, diagram_analysis, console=console,
        )

    # Stage 4: Merge metadata
    if merge_metadata:
        latex = merge_metadata_into_latex(latex, primary, diagram_analysis, difficulty_result)
        console.print("[green]✓[/green] Metadata merged")

    # Stage 5: Ideas, Alternates, Variants
    problem, solution = extract_problem_solution(latex)

    ideas = None
    if generate_ideas and problem and solution:
        ideas = extract_ideas_stage(
            problem, solution, cache=cache, problem_id=problem_id, console=console,
        )

    alternate_solutions = []
    if generate_alternate and problem and solution:
        alternate_solutions = generate_alternate_stage(
            problem, solution, ideas, cache=cache, problem_id=problem_id, console=console,
        )

    final_classification = ClassificationResult(
        subject=primary.subject,
        question_type=primary.question_type,
        has_diagram=primary.has_diagram,
        confidence=primary.confidence,
        classified_from=primary.classified_from,
    )

    variants = generate_variants_stage(
        latex, variant_types, ideas, final_classification,
        use_context=use_context, cache=cache, problem_id=problem_id, console=console,
    )

    return PipelineResult(
        source_path=image_path,
        classification=final_classification,
        latex=latex,
        tikz_code=tikz_code,
        ideas=ideas,
        alternate_solutions=alternate_solutions,
        variants=variants,
    )
