from typing import List, Tuple
import os
import time

from .data import (
    ObjectInfo, FunctionInfo, FunctionImplError,
    ProjectSynthMetadata, FunctionImpl, UnharnessableFunction,
)
from ..data.metadata import UsageTracker
from ..data.project_info import ProjectInfo
from .data import ProjectMetadata

from .scan_objects import scan_objects
from .scan_functions import scan_functions, scan_function_names
from .synth_functions import incremental_synthesis_loop
from .data_classifier import classify_data
from .fill import run_fill_batch, analyze_blocking_missing_types
from .validation import run_validate_object
from .build_generators import build_generator

from . import ui

SYNTH_RETRIES = 3


def process_header_functions(
    header_path: str,
    objects: List[ObjectInfo],
    info: ProjectInfo,
    threads: int,
    usage: UsageTracker,
) -> Tuple[List[FunctionInfo], List[FunctionImpl], List[FunctionImplError], List[UnharnessableFunction]]:
    code = open(header_path, 'rb').read().decode('utf-8', errors='ignore')

    function_names = scan_function_names(code, usage)
    functions = scan_functions(code, function_names, usage, quiet=True)

    valid_stubs, failed_stubs, unharnessable = incremental_synthesis_loop(
        objects, functions, usage, threads, info,
        show_progress=False,
    )

    return functions, valid_stubs, failed_stubs, unharnessable


def main(args):
    info = ProjectInfo.model_validate_json(open(args.info).read())

    merged_header = ''
    for header in args.headers:
        merged_header += f'<CODE FILE="{header}">\n'
        with open(header, 'rb') as f:
            merged_header += f.read().decode('utf-8', errors='ignore')
        merged_header += '</CODE>\n'

    usage = UsageTracker()
    start_time = time.time()

    ui.pipeline_header()

    ui.phase("scan-objects", "Identifying objects in header files.")

    objects = []

    header_names = ", ".join(os.path.basename(h) for h in args.headers)

    if len(args.headers) <= 3:
        s_objects = time.time()
        tp = ui.TokenProgress()
        with ui.spinner(f"Scanning [bold]{header_names}[/bold] for objects", token_progress=tp):
            objects = scan_objects(merged_header, usage, token_progress=tp)
        t_objects = time.time() - s_objects
    else:
        s_objects = time.time()
        objects = ui.parallel_run(
            args.threads, scan_objects,
            [(open(h, 'rb').read().decode('utf-8', errors='ignore'), usage) for h in args.headers],
            item_labels=[os.path.basename(h) for h in args.headers],
        )
        t_objects = time.time() - s_objects

    ui.log('SCAN', f"Found [cyan]{len(objects)}[/cyan] objects")

    s_validate_objects = time.time()
    objects: List[ObjectInfo] = ui.parallel_run(
        args.threads, run_validate_object,
        [(info, o) for o in objects],
        item_labels=[o.ptr_type for o in objects],
        ok=lambda r: len(r) > 0,
    )
    t_validate_objects = time.time() - s_validate_objects

    ui.phase_result(
        f"[cyan]{len(objects)}[/cyan] valid objects",
        t_objects + t_validate_objects,
    )

    functions = []
    valid_stubs = []
    failed_stubs = []
    unharnessable_functions: List[UnharnessableFunction] = []

    t_functions = 0
    t_synthesis = 0

    if len(args.headers) <= 3:
        ui.phase("scan-names", "Discovering function names in source code.")

        s_functions = time.time()
        tp = ui.TokenProgress()
        with ui.spinner(f"Scanning [bold]{header_names}[/bold] for function names", token_progress=tp):
            function_names = scan_function_names(merged_header, usage, token_progress=tp)
        ui.log('SCAN', f"Found [cyan]{len(function_names)}[/cyan] function names")

        ui.phase("scan-funcs", "Analyzing function signatures and parameters.")

        tp = ui.TokenProgress()
        with ui.spinner(f"Scanning [bold]{header_names}[/bold] for function signatures", token_progress=tp):
            functions = scan_functions(merged_header, function_names, usage, quiet=True, token_progress=tp)
        t_functions = time.time() - s_functions
        ui.phase_result(
            f"[cyan]{len(function_names)}[/cyan] names  ·  "
            f"[cyan]{len(functions)}[/cyan] functions",
            t_functions,
        )

        ui.phase("synth-functions", "Synthesizing function implementations.")

        s_synthesis = time.time()
        valid_stubs, failed_stubs, unharnessable = incremental_synthesis_loop(
            objects, functions, usage, args.threads, info,
            verbose=args.verbose,
            token_progress=ui.TokenProgress(),
        )
        unharnessable_functions.extend(unharnessable)
        t_synthesis = time.time() - s_synthesis

        parts = [f"[green]{len(valid_stubs)}[/green] synthesized"]
        if failed_stubs:
            parts.append(f"[red]{len(failed_stubs)}[/red] failed")
        if unharnessable:
            parts.append(f"[yellow]{len(unharnessable)}[/yellow] unharnessable")
        ui.phase_result("  ·  ".join(parts), t_synthesis)

    else:
        ui.phase("scan-names", "Discovering function names in source code.")
        ui.phase("scan-funcs", "Analyzing function signatures and parameters.")
        ui.info("Large project — independent header processing")

        s_functions = time.time()
        res = ui.parallel_run_raw(
            args.threads, process_header_functions,
            [(hp, objects, info, args.threads, usage) for hp in args.headers],
            item_labels=[os.path.basename(hp) for hp in args.headers],
        )
        for p_functions, p_valid_stubs, p_failed_stubs, p_unharnessable in res:
            functions.extend(p_functions)
            valid_stubs.extend(p_valid_stubs)
            failed_stubs.extend(p_failed_stubs)
            unharnessable_functions.extend(p_unharnessable)
        t_synthesis = time.time() - s_functions

        ui.phase_result(
            f"[cyan]{len(functions)}[/cyan] functions scanned",
            t_synthesis,
        )

        ui.phase("synth-functions", "Synthesizing function implementations.")

        parts = [f"[green]{len(valid_stubs)}[/green] synthesized"]
        if failed_stubs:
            parts.append(f"[red]{len(failed_stubs)}[/red] failed")
        if unharnessable_functions:
            parts.append(f"[yellow]{len(unharnessable_functions)}[/yellow] unharnessable")
        ui.phase_result("  ·  ".join(parts), t_synthesis)

    meta = ProjectMetadata(
        objects=objects,
        functions=functions,
        endpoints=valid_stubs,
        unharnessable_functions=unharnessable_functions,
        usage=usage,
    )

    synth_meta = ProjectSynthMetadata(
        endpoint_failures=failed_stubs,
    )

    selective_context = len(meta.endpoints) > 300 or len(meta.objects) > 300

    ui.phase("synth-fill", "Filling missing types for constructability.")

    with ui.spinner("Analyzing constructability"):
        target_types = analyze_blocking_missing_types(valid_stubs, objects, max_depth=6)

    t_fill = 0
    if len(target_types) > 0:
        ui.log('FILL', f"Found [yellow]{len(target_types)}[/yellow] blocking types to fill")

        s_fill = time.time()
        valid_stubs, failed_stubs = run_fill_batch(
            info, meta, target_types, usage, args.threads, selective_context,
            verbose=args.verbose,
            token_progress=ui.TokenProgress(),
        )
        meta.endpoints.extend(valid_stubs)
        synth_meta.endpoint_failures.extend(failed_stubs)
        t_fill = time.time() - s_fill

        parts = [f"[green]{len(valid_stubs)}[/green] filled"]
        if failed_stubs:
            parts.append(f"[red]{len(failed_stubs)}[/red] failed")
        ui.phase_result("  ·  ".join(parts), t_fill)
    else:
        ui.success("All types constructible — nothing to fill")

    ui.phase("classify-data", "Classifying data types and endpoints.")

    s_classify = time.time()
    data_types, classifiers = classify_data(
        meta.objects, meta.functions, meta.endpoints, usage, selective_context,
        token_progress=ui.TokenProgress(),
    )
    t_classify = time.time() - s_classify
    meta.data_types = data_types
    meta.endpoint_classifiers = classifiers

    ui.phase_result(
        f"[cyan]{len(data_types)}[/cyan] data types  ·  "
        f"[cyan]{len(classifiers)}[/cyan] classifiers",
        t_classify,
    )

    ui.phase("build-generators", "Constructing data generators.")

    s_build_generators = time.time()
    generators = ui.parallel_run(
        args.threads, build_generator,
        [
            (d, meta.objects, meta.functions, meta.endpoints,
             meta.endpoint_classifiers, usage, selective_context, True)
            for d in meta.data_types
        ],
        item_labels=[d.data_type for d in meta.data_types],
        ok=lambda r: len(r) > 0,
    )
    t_build_generators = time.time() - s_build_generators
    meta.generators = generators

    ui.phase_result(
        f"[green]{len(meta.generators)}[/green] generators built",
        t_build_generators,
    )

    # ── Summary ──────────────────────────────────────────────────────────
    end_time = time.time()
    meta.metadata['total_time'] = end_time - start_time
    meta.metadata['scan_objects'] = t_objects
    meta.metadata['validate_objects'] = t_validate_objects
    meta.metadata['scan_functions'] = t_functions
    meta.metadata['synthesis'] = t_synthesis
    meta.metadata['fill'] = t_fill
    meta.metadata['classify'] = t_classify
    meta.metadata['build_generators'] = t_build_generators

    ui.pipeline_summary(meta, end_time - start_time, usage)

    open(args.output, 'w').write(meta.model_dump_json(indent=2))
    open(args.output.replace('.json', '_failed.json'), 'w').write(synth_meta.model_dump_json(indent=2))


def register(subparsers):
    scan_parser = subparsers.add_parser('run')
    scan_parser.add_argument('--info', type=str, default='info.json', help='Path to the project info file')
    scan_parser.add_argument('--headers', type=str, required=True, nargs='+', help='List of header files to scan')
    scan_parser.add_argument('--threads', type=int, default=1, help='Number of threads to use')
    scan_parser.add_argument('-o', '--output', type=str, required=True)
    scan_parser.add_argument('--verbose', action='store_true', help='Verbose output')
    scan_parser.set_defaults(func=main)
