from __future__ import annotations

import re
from collections.abc import Iterable
from typing import Any
from uuid import uuid4

from skills_learning.models import (
    LOCAL_SYNTHESIS_SOURCE,
    FlowStep,
    FlowTemplateSkill,
    SkillMetadata,
)
from skills_learning.store import SkillStore

READ_ONLY_TOOL_NAMES = {
    "describe_screen",
    "describe_screen_semantic",
    "screenshot",
    "take_screenshot",
    "ui_tree",
    "get_ui_tree",
    "find_text",
    "find_text_on_screen",
    "find_icon",
    "find_icon_on_screen",
}


def extract_tool_calls(
    task_id: str,
    timeline: Any,
) -> list[FlowStep]:
    records = _timeline_records(timeline, task_id)
    steps: list[FlowStep] = []
    for record in records:
        tool_call = _record_value(record, "tool_call")
        if not isinstance(tool_call, dict):
            continue
        tool_name = str(tool_call.get("action") or tool_call.get("tool_name") or "")
        if not tool_name or tool_name in READ_ONLY_TOOL_NAMES:
            continue
        steps.append(
            FlowStep(
                tool_name=tool_name,
                args=dict(tool_call.get("args") or {}),
                purpose=_optional_text(tool_call.get("purpose")),
                expected_outcome=_optional_text(tool_call.get("expected_outcome")),
            )
        )
    return steps


def find_matching_skeleton(
    steps: list[FlowStep],
    store: SkillStore | None,
) -> FlowTemplateSkill | None:
    if store is None:
        return None
    sequence = _tool_sequence(steps)
    for skill in store.list_latest():
        if _tool_sequence(skill.steps) == sequence:
            return skill
    return None


def synthesize_flow_skill(
    goal: str,
    timeline: Any,
    *,
    task_id: str = "",
    store: SkillStore | None = None,
) -> FlowTemplateSkill:
    records = _timeline_records(timeline, task_id)
    executed_steps = extract_tool_calls(task_id, records)
    matched = find_matching_skeleton(executed_steps, store)

    if matched is None:
        steps = executed_steps
        parameters: dict[str, dict[str, Any]] = {}
        name = _skill_name_from_goal(goal)
    else:
        steps, parameters = promote_parameters(
            matched.steps,
            executed_steps,
            records=records,
            existing_parameters=matched.parameters,
        )
        name = matched.name

    return FlowTemplateSkill(
        metadata=SkillMetadata(
            id=uuid4().hex,
            name=name,
            description=f"Learned flow for: {goal}",
            kind="flow_template",
            source=LOCAL_SYNTHESIS_SOURCE,
            version=matched.version if matched else 1,
            parent_version_id=matched.parent_version_id if matched else None,
            originating_goal=goal,
        ),
        steps=steps,
        parameters=parameters,
    )


def promote_parameters(
    stored_steps: list[FlowStep],
    executed_steps: list[FlowStep],
    *,
    records: Iterable[Any] = (),
    existing_parameters: dict[str, dict[str, Any]] | None = None,
) -> tuple[list[FlowStep], dict[str, dict[str, Any]]]:
    parameters = {
        name: dict(schema) for name, schema in (existing_parameters or {}).items()
    }
    parameterized_steps = [
        FlowStep(
            step.tool_name,
            dict(step.args),
            purpose=step.purpose,
            expected_outcome=step.expected_outcome,
        )
        for step in executed_steps
    ]
    used_names = set(parameters)
    parameter_index = len(used_names) + 1

    for step_index, (stored_step, executed_step) in enumerate(
        zip(stored_steps, executed_steps, strict=False)
    ):
        for arg_name, executed_value in executed_step.args.items():
            stored_value = stored_step.args.get(arg_name)
            if stored_value == executed_value:
                continue
            if _is_placeholder(stored_value):
                parameterized_steps[step_index].args[arg_name] = stored_value
                continue
            if _is_placeholder(executed_value):
                continue

            preferred_name = _parameter_name_from_semantic_record(
                list(records),
                step_index,
            )
            parameter_name = _unique_parameter_name(
                preferred_name or f"param_{parameter_index}",
                used_names,
            )
            used_names.add(parameter_name)
            parameter_index += 1
            parameterized_steps[step_index].args[arg_name] = f"{{{parameter_name}}}"
            parameters[parameter_name] = _parameter_schema(
                arg_name,
                executed_value,
                preferred_name is not None,
            )

    return parameterized_steps, parameters


def _timeline_records(timeline: Any, task_id: str) -> list[Any]:
    if isinstance(timeline, list):
        return list(timeline)
    if hasattr(timeline, "read"):
        return list(timeline.read(task_id))
    return list(timeline)


def _record_value(record: Any, key: str) -> Any:
    if isinstance(record, dict):
        return record.get(key)
    return getattr(record, key, None)


def _optional_text(value: Any) -> str | None:
    if not isinstance(value, str):
        return None
    return value.strip() or None


def _tool_sequence(steps: list[FlowStep]) -> list[str]:
    return [step.tool_name for step in steps]


def _is_placeholder(value: Any) -> bool:
    return isinstance(value, str) and value.startswith("{") and value.endswith("}")


def _parameter_schema(
    arg_name: str,
    value: Any,
    from_semantic_label: bool,
) -> dict[str, Any]:
    return {
        "type": _json_type(value),
        "description": (
            f"Value for {arg_name} inferred from a semantic widget label"
            if from_semantic_label
            else f"Value for {arg_name}"
        ),
    }


def _json_type(value: Any) -> str:
    if isinstance(value, bool):
        return "boolean"
    if isinstance(value, int | float):
        return "number"
    if isinstance(value, list):
        return "array"
    if isinstance(value, dict):
        return "object"
    return "string"


def _parameter_name_from_semantic_record(
    records: list[Any],
    step_index: int,
) -> str | None:
    if step_index >= len(records):
        return None
    result = _record_value(records[step_index], "result")
    semantic_scene = _semantic_scene_payload(result)
    if not isinstance(semantic_scene, dict):
        return None
    widgets = semantic_scene.get("widgets")
    if not isinstance(widgets, list):
        return None
    for widget in widgets:
        if isinstance(widget, dict) and widget.get("purpose"):
            return _sanitize_name(str(widget["purpose"]))
    return None


def _semantic_scene_payload(result: Any) -> Any:
    if not isinstance(result, dict):
        return None
    if "semantic_scene" in result:
        return result["semantic_scene"]
    nested = result.get("result")
    if isinstance(nested, dict):
        return nested.get("semantic_scene")
    return None


def _unique_parameter_name(name: str, used_names: set[str]) -> str:
    candidate = _sanitize_name(name) or "param"
    if candidate not in used_names:
        return candidate
    index = 2
    while f"{candidate}_{index}" in used_names:
        index += 1
    return f"{candidate}_{index}"


def _sanitize_name(value: str) -> str:
    sanitized = re.sub(r"[^0-9a-zA-Z]+", "_", value.strip().lower()).strip("_")
    if sanitized and sanitized[0].isdigit():
        return f"param_{sanitized}"
    return sanitized


def _skill_name_from_goal(goal: str) -> str:
    return goal.strip() or "learned flow"