"""Integrator — Composition + I/O tracing workflow.

Integration is itself a black box with a contract. The parent component
has its own ComponentContract and ContractTestSuite. Integration writes
glue code wiring children together.

When parent-level tests fail, I/O tracing finds the failure point.
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Callable
from datetime import datetime

from pact.agents.base import AgentBase
from pact.project import ProjectManager
from pact.schemas import (
    ComponentContract,
    ContractTestSuite,
    DecompositionTree,
    TestResults,
)
from pact.test_harness import run_contract_tests

logger = logging.getLogger(__name__)

GLUE_SYSTEM = """You are starting fresh on this integration task with no prior context.

You are an integration engineer wiring child components into the parent's
interface. Glue code handles data transformation and routing between
components but adds no business logic. All parent functions delegate
to children. Error propagation must match the parent contract."""

GLUE_SYSTEM_TS = """You are starting fresh on this integration task with no prior context.

You are an integration engineer producing TypeScript glue code that
composes child implementations into the parent's interface. Import
children using ESM imports, export using named exports. Glue code
handles data transformation and routing — no business logic. Use
strict mode, unknown instead of any."""

GLUE_SYSTEM_JS = """You are starting fresh on this integration task with no prior context.

You are an integration engineer producing JavaScript glue code that
composes child implementations into the parent's interface. Import
children using ESM imports with .js extensions, export using named
exports. Glue code handles data transformation and routing — no
business logic. Plain JavaScript ES6+ modules only."""


async def integrate_component(
    agent: AgentBase,
    project: ProjectManager,
    parent_id: str,
    parent_contract: ComponentContract,
    parent_test_suite: ContractTestSuite,
    child_contracts: dict[str, ComponentContract],
    max_attempts: int = 3,
    sops: str = "",
) -> TestResults:
    """Integrate child components into a parent.

    Returns:
        TestResults from running parent-level tests.
    """
    from pydantic import BaseModel

    language = project.language
    is_ts = language == "typescript"
    is_js = language == "javascript"
    glue_ext = ".js" if is_js else (".ts" if is_ts else ".py")

    class GlueResponse(BaseModel):
        """Generated glue code."""
        glue_code: str
        composition_test: str = ""

    children_summary = "\n".join(
        f"  - {cid}: {c.name} — {', '.join(f.name for f in c.functions)}"
        for cid, c in child_contracts.items()
    )

    parent_funcs = "\n".join(
        f"  - {f.name}({', '.join(i.name + ': ' + i.type_ref for i in f.inputs)}) -> {f.output_type}"
        for f in parent_contract.functions
    )

    prior_failures: list[str] = []

    for attempt in range(1, max_attempts + 1):
        failure_context = ""
        if prior_failures:
            failure_context = (
                "\nPrior failures:\n"
                + "\n".join(f"  - {f}" for f in prior_failures)
            )

        if is_ts:
            lang_label = "TypeScript"
            import_hint = (
                "- Import from each child using ESM imports "
                "(e.g., `import { fn } from './child_module';`)"
            )
        elif is_js:
            lang_label = "JavaScript"
            import_hint = (
                "- Import from each child using ESM imports with .js extensions "
                "(e.g., `import { fn } from './child_module.js';`)"
            )
        else:
            lang_label = "Python"
            import_hint = "- Import from each child's module"

        # Build full child contract JSON
        child_contracts_json = "\n\n".join(
            f'"{cid}":\n{c.model_dump_json(indent=2)}'
            for cid, c in child_contracts.items()
        )

        # Load child implementation source code
        child_impls = ""
        for cid in child_contracts:
            impl_src = project.impl_src_dir(cid)
            if impl_src.exists():
                for src_file in impl_src.rglob("*"):
                    if src_file.is_file() and src_file.suffix in (".py", ".ts", ".js"):
                        child_impls += f"\n\n=== {cid} implementation ({src_file.name}) ===\n{src_file.read_text()}"

        prompt = f"""Generate glue code to compose children into the parent interface.

Parent: {parent_contract.name} (id: {parent_id})
Parent functions:
{parent_funcs}

Children:
{children_summary}

Parent contract (JSON):
{parent_contract.model_dump_json(indent=2)}

Child contracts (full JSON):
{child_contracts_json}
{f'{chr(10)}Child implementations:{child_impls}' if child_impls else ''}
{failure_context}

Generate:
1. glue_code: {lang_label} module that imports children and implements parent interface
2. composition_test: Optional additional integration tests

The glue code should:
{import_hint}
- Implement each parent function by delegating to appropriate children
- Handle data transformation between child interfaces
- Propagate errors according to parent contract"""

        system = GLUE_SYSTEM_TS if is_ts else (GLUE_SYSTEM_JS if is_js else GLUE_SYSTEM)
        response, _, _ = await agent.assess(GlueResponse, prompt, system)

        # Save glue code
        comp_dir = project.composition_dir(parent_id)
        glue_path = comp_dir / f"glue{glue_ext}"
        glue_path.write_text(response.glue_code)

        if response.composition_test:
            test_ext = ".test.ts" if is_ts else ".py"
            # Save composition tests to visible tests dir
            visible_test_dir = project._visible_tests_dir / parent_id
            visible_test_dir.mkdir(parents=True, exist_ok=True)
            test_path = visible_test_dir / f"composition_test{test_ext}"
            test_path.write_text(response.composition_test)

        project.append_audit(
            "integration",
            f"{parent_id} attempt {attempt}",
        )

        # Run parent tests
        test_file = project.test_code_path(parent_id)
        if not test_file.exists() and parent_test_suite.generated_code:
            test_file.parent.mkdir(parents=True, exist_ok=True)
            test_file.write_text(parent_test_suite.generated_code)

        # Include child implementation src/ dirs so glue code can import them
        child_paths = [
            project.impl_src_dir(cid) for cid in child_contracts
        ]

        test_results = await run_contract_tests(
            test_file, comp_dir, extra_paths=child_paths,
            language=language,
            project_dir=project.project_dir,
        )

        # Save results to internal composition dir
        internal_comp = project._internal_composition_dir(parent_id)
        results_path = internal_comp / "test_results.json"
        results_path.write_text(test_results.model_dump_json(indent=2))

        if test_results.all_passed:
            logger.info(
                "Integration %s passed all %d tests on attempt %d",
                parent_id, test_results.total, attempt,
            )
            return test_results

        for failure in test_results.failure_details:
            prior_failures.append(
                f"Test '{failure.test_id}': {failure.error_message}"
            )

        logger.warning(
            "Integration %s failed %d/%d tests on attempt %d",
            parent_id, test_results.failed + test_results.errors,
            test_results.total, attempt,
        )

    return test_results


async def integrate_all(
    agent: AgentBase,
    project: ProjectManager,
    tree: DecompositionTree,
    max_attempts: int = 3,
    sops: str = "",
    parallel: bool = False,
    max_concurrent: int = 4,
    agent_factory: Callable[[], AgentBase] | None = None,
) -> dict[str, TestResults]:
    """Integrate all non-leaf components, deepest first.

    When parallel=True, non-leaves at the same depth are integrated
    concurrently (they're independent since their children are already done).
    Groups execute in order: deepest first, so children finish before parents.

    Returns:
        Dict of parent_id -> TestResults.
    """
    contracts = project.load_all_contracts()
    test_suites = project.load_all_test_suites()
    results: dict[str, TestResults] = {}

    # Get depth-ordered groups (deepest first)
    if parallel:
        groups = tree.non_leaf_parallel_groups()
    else:
        # Sequential: use topological order, non-leaves only
        order = tree.topological_order()
        groups = [[cid] for cid in order
                  if tree.nodes.get(cid) and tree.nodes[cid].children]

    async def _integrate_one(component_id: str) -> tuple[str, TestResults] | None:
        node = tree.nodes.get(component_id)
        if not node or not node.children:
            return None

        if component_id not in contracts:
            logger.warning("No contract for parent %s", component_id)
            return None

        child_contracts = {
            cid: contracts[cid]
            for cid in node.children
            if cid in contracts
        }

        test_suite = test_suites.get(component_id)
        if not test_suite:
            logger.warning("No test suite for parent %s", component_id)
            return None

        int_agent = agent_factory() if (parallel and agent_factory) else agent
        try:
            test_results = await integrate_component(
                int_agent, project, component_id,
                contracts[component_id],
                test_suite,
                child_contracts,
                max_attempts=max_attempts,
                sops=sops,
            )
        finally:
            if parallel and agent_factory and int_agent is not agent:
                await int_agent.close()

        node.implementation_status = (
            "tested" if test_results.all_passed else "failed"
        )
        node.test_results = test_results
        return component_id, test_results

    for group in groups:
        if parallel and len(group) > 1:
            sem = asyncio.Semaphore(max_concurrent)

            async def _limited(cid: str) -> tuple[str, TestResults] | None:
                async with sem:
                    return await _integrate_one(cid)

            gather_results = await asyncio.gather(*[_limited(cid) for cid in group])
            for result in gather_results:
                if result:
                    results[result[0]] = result[1]
        else:
            for component_id in group:
                result = await _integrate_one(component_id)
                if result:
                    results[result[0]] = result[1]

    project.save_tree(tree)
    return results


async def integrate_component_iterative(
    project: ProjectManager,
    parent_id: str,
    parent_contract: ComponentContract,
    parent_test_suite: ContractTestSuite,
    child_contracts: dict[str, ComponentContract],
    budget: object,  # BudgetTracker
    model: str = "claude-opus-4-6",
    sops: str = "",
    external_context: str = "",
    learnings: str = "",
    max_turns: int = 30,
    timeout: int = 600,
) -> TestResults:
    """Integrate child components via iterative Claude Code (write glue -> test -> fix).

    Instead of asking the API to produce a JSON blob of glue code, gives Claude Code
    full tool access to read child implementations, write glue code, run parent tests,
    read errors, and iterate within a single session.

    Returns:
        TestResults from running parent-level tests.
    """
    from pact.backends.claude_code import ClaudeCodeBackend

    language = project.language
    is_ts = language == "typescript"
    is_js = language == "javascript"
    file_ext = ".js" if is_js else (".ts" if is_ts else ".py")

    children_summary = "\n".join(
        f"  - {cid}: {c.name} — {', '.join(f.name for f in c.functions)}"
        for cid, c in child_contracts.items()
    )

    parent_funcs = "\n".join(
        f"  - {f.name}({', '.join(i.name + ': ' + i.type_ref for i in f.inputs)}) -> {f.output_type}"
        for f in parent_contract.functions
    )

    # Gather child implementation paths (absolute for env vars, relative for prompts)
    child_src_dirs = {
        cid: project.impl_src_dir(cid) for cid in child_contracts
    }
    child_src_dirs_rel = {
        cid: path.relative_to(project.project_dir)
        for cid, path in child_src_dirs.items()
    }
    child_impl_listing = "\n".join(
        f"  - {cid}: {path}/"
        for cid, path in child_src_dirs_rel.items()
    )

    # Write test file so the agent can run it
    test_file = project.test_code_path(parent_id)
    if not test_file.exists() and parent_test_suite.generated_code:
        test_file.parent.mkdir(parents=True, exist_ok=True)
        test_file.write_text(parent_test_suite.generated_code)

    comp_dir = project.composition_dir(parent_id)
    comp_dir_rel = comp_dir.relative_to(project.project_dir)
    test_file_rel = test_file.relative_to(project.project_dir)

    module_name = parent_id.replace("-", "_")

    if is_ts or is_js:
        # Build NODE_PATH for TypeScript/JavaScript module resolution
        node_path_parts = [str(comp_dir), str(comp_dir.parent)]
        node_path_parts.extend(str(p) for p in child_src_dirs.values())
        env_path_str = ":".join(node_path_parts)

        ts_specific = ""
        if is_ts:
            ts_specific = "   - Use named exports only (no default exports)\n"
        js_specific = ""
        if is_js:
            js_specific = (
                "   - Use ESM imports with .js file extensions\n"
                "   - Do NOT use TypeScript syntax — plain JavaScript only\n"
            )

        prompt = f"""You are an integration engineer. Wire child components together into the parent interface.

## Parent Component: {parent_contract.name} (id: {parent_id})

Parent functions:
{parent_funcs}

Parent contract (JSON):
{parent_contract.model_dump_json(indent=2)}

## Children

{children_summary}

Child contracts:
{chr(10).join(f'{cid}: {c.model_dump_json(indent=2)}' for cid, c in child_contracts.items())}

## Child Implementation Locations

{child_impl_listing}

## CRITICAL: Module Structure Convention

The test file imports from: `./src/{module_name}`

You MUST write your glue module at: {comp_dir_rel}/src/{module_name}{file_ext}

Create the directory if needed: mkdir -p {comp_dir_rel}/src/

Children are importable using ESM imports. For example:
{chr(10).join(f'  import {{ ... }} from "./{cid.replace("-", "_")}";' for cid in child_contracts)}

## Your Task

1. Read each child implementation to understand their actual APIs:
{chr(10).join(f'   - {path}/{cid.replace("-", "_")}{file_ext}' for cid, path in child_src_dirs_rel.items())}
2. Read the parent test file: {test_file_rel}
3. Write your glue module at: {comp_dir_rel}/src/{module_name}{file_ext}
   - Import from each child module using ESM imports
   - Re-export ALL types and functions that the test file imports
   - Implement each parent function by delegating to appropriate children
   - Match the exact type names, function signatures, and enum values from the contract
{ts_specific}{js_specific}   - Use named exports only (no default exports)
   - Do NOT add business logic — only wiring and delegation
4. Run tests:
   NODE_PATH="{env_path_str}" npx vitest run {test_file_rel}
5. If tests fail, read the errors, fix your glue code, and re-run
6. Keep iterating until ALL tests pass

{f'SOPs: {sops}' if sops else ''}
{f'Context: {external_context}' if external_context else ''}
{f'Learnings: {learnings}' if learnings else ''}
"""
    else:
        # Build PYTHONPATH that includes child implementation dirs
        pythonpath_parts = [str(comp_dir), str(comp_dir.parent)]
        pythonpath_parts.extend(str(p) for p in child_src_dirs.values())
        pythonpath_str = ":".join(pythonpath_parts)

        prompt = f"""You are an integration engineer. Wire child components together into the parent interface.

## Parent Component: {parent_contract.name} (id: {parent_id})

Parent functions:
{parent_funcs}

Parent contract (JSON):
{parent_contract.model_dump_json(indent=2)}

## Children

{children_summary}

Child contracts:
{chr(10).join(f'{cid}: {c.model_dump_json(indent=2)}' for cid, c in child_contracts.items())}

## Child Implementation Locations

{child_impl_listing}

## CRITICAL: Module Structure Convention

The test file imports: `from {module_name} import ...`

You MUST write your glue module at: {comp_dir_rel}/src/{module_name}.py

Create the directory if needed: mkdir -p {comp_dir_rel}/src/

Children are importable directly by name (they're on PYTHONPATH). For example:
{chr(10).join(f'  import {cid.replace("-", "_")}' for cid in child_contracts)}

Do NOT use sys.path manipulation. Just import children by their module name.

## Your Task

1. Read each child implementation to understand their actual APIs:
{chr(10).join(f'   - {path}/{cid.replace("-", "_")}.py' for cid, path in child_src_dirs_rel.items())}
2. Read the parent test file: {test_file_rel}
3. Write your glue module at: {comp_dir_rel}/src/{module_name}.py
   - Create {comp_dir_rel}/src/__init__.py if needed
   - Import from each child module by name (e.g. `import <child_module_name>`)
   - Re-export ALL types and functions that the test file imports
   - Implement each parent function by delegating to appropriate children
   - Match the exact type names, function signatures, and enum values from the contract
   - Do NOT add business logic — only wiring and delegation
4. Run tests with correct PYTHONPATH:
   PYTHONPATH="{pythonpath_str}" python3 -m pytest {test_file_rel} -v
5. If tests fail, read the errors, fix your glue code, and re-run
6. Keep iterating until ALL tests pass

{f'SOPs: {sops}' if sops else ''}
{f'Context: {external_context}' if external_context else ''}
{f'Learnings: {learnings}' if learnings else ''}
"""

    logger.info("Integrating %s iteratively via Claude Code (%s)", parent_id, model)

    backend = ClaudeCodeBackend(
        budget=budget,
        model=model,
        repo_path=project.project_dir,
        timeout=timeout,
    )

    try:
        await backend.implement(
            prompt=prompt,
            working_dir=project.project_dir,
            max_turns=max_turns,
            timeout=timeout,
        )
    except Exception as e:
        logger.error("Iterative integration failed for %s: %s", parent_id, e)

    project.append_audit(
        "integration",
        f"{parent_id} iterative claude_code ({model})",
    )

    # Run parent tests for official results — include child src/ dirs
    child_paths = [
        project.impl_src_dir(cid) for cid in child_contracts
    ]
    test_results = await run_contract_tests(
        test_file, comp_dir, extra_paths=child_paths,
        language=language,
        project_dir=project.project_dir,
    )

    # Save results to internal composition dir
    internal_comp = project._internal_composition_dir(parent_id)
    results_path = internal_comp / "test_results.json"
    results_path.write_text(test_results.model_dump_json(indent=2))

    project.append_audit(
        "test_run",
        f"integration {parent_id}: {test_results.passed}/{test_results.total} passed",
    )

    logger.info(
        "Integration %s iterative result: %d/%d passed",
        parent_id, test_results.passed, test_results.total,
    )

    return test_results


async def integrate_all_iterative(
    project: ProjectManager,
    tree: DecompositionTree,
    budget: object,  # BudgetTracker
    model: str = "claude-opus-4-6",
    sops: str = "",
    parallel: bool = False,
    max_concurrent: int = 4,
    external_context: str = "",
    learnings: str = "",
    max_turns: int = 30,
    timeout: int = 600,
) -> dict[str, TestResults]:
    """Integrate all non-leaf components via iterative Claude Code, deepest first.

    Returns:
        Dict of parent_id -> TestResults.
    """
    contracts = project.load_all_contracts()
    test_suites = project.load_all_test_suites()
    results: dict[str, TestResults] = {}

    # Get depth-ordered groups (deepest first)
    if parallel:
        groups = tree.non_leaf_parallel_groups()
    else:
        order = tree.topological_order()
        groups = [[cid] for cid in order
                  if tree.nodes.get(cid) and tree.nodes[cid].children]

    async def _integrate_one(component_id: str) -> tuple[str, TestResults] | None:
        node = tree.nodes.get(component_id)
        if not node or not node.children:
            return None

        if component_id not in contracts:
            logger.warning("No contract for parent %s", component_id)
            return None

        child_contracts = {
            cid: contracts[cid]
            for cid in node.children
            if cid in contracts
        }

        test_suite = test_suites.get(component_id)
        if not test_suite:
            logger.warning("No test suite for parent %s", component_id)
            return None

        test_results = await integrate_component_iterative(
            project, component_id,
            contracts[component_id],
            test_suite,
            child_contracts,
            budget=budget,
            model=model,
            sops=sops,
            external_context=external_context,
            learnings=learnings,
            max_turns=max_turns,
            timeout=timeout,
        )

        node.implementation_status = (
            "tested" if test_results.all_passed else "failed"
        )
        node.test_results = test_results
        return component_id, test_results

    for group in groups:
        if parallel and len(group) > 1:
            sem = asyncio.Semaphore(max_concurrent)

            async def _limited(cid: str) -> tuple[str, TestResults] | None:
                async with sem:
                    return await _integrate_one(cid)

            gather_results = await asyncio.gather(*[_limited(cid) for cid in group])
            for result in gather_results:
                if result:
                    results[result[0]] = result[1]
        else:
            for component_id in group:
                result = await _integrate_one(component_id)
                if result:
                    results[result[0]] = result[1]

    project.save_tree(tree)
    return results
