import re
from typing import List, Set, Dict, Any
from ..schemas.agent_spec import AgentSpec
from ..schemas.conversation import ValidationError, ValidationResult

class WorkflowValidator:
    def validate(self, spec: AgentSpec) -> ValidationResult:
        errors = []
        steps_dict = {s.id: s for s in spec.steps}
        step_ids = list(steps_dict.keys())
        
        # 1. Duplicate step IDs
        all_ids = [s.id for s in spec.steps]
        if len(all_ids) != len(set(all_ids)):
            errors.append(ValidationError(
                code="DUPLICATE_STEP_ID",
                message="Duplicate step IDs found in the workflow",
                path="steps"
            ))

            
        # 2. Depends_on exists
        for step in spec.steps:
            for dep in step.depends_on:
                if dep not in steps_dict:
                    errors.append(ValidationError(
                        code="MISSING_DEPENDENCY",
                        message=f"Step '{step.id}' depends on missing step '{dep}'",
                        path=f"steps.{step.id}.depends_on"
                    ))

        # 3. Circular dependency detection (Cycle Detection using DFS)
        visited = {}  # state: 0=unvisited, 1=visiting, 2=visited
        for s_id in step_ids:
            visited[s_id] = 0

        def dfs(s_id: str, path: List[str]) -> bool:
            visited[s_id] = 1  # visiting
            step = steps_dict[s_id]
            for dep in step.depends_on:
                if dep not in steps_dict:
                    continue
                if visited[dep] == 1:
                    cycle_path = " -> ".join(path + [dep])
                    errors.append(ValidationError(
                        code="CIRCULAR_DEPENDENCY",
                        message=f"Circular dependency detected: {cycle_path}",
                        path=f"steps.{s_id}.depends_on"
                    ))
                    return True
                elif visited[dep] == 0:
                    if dfs(dep, path + [dep]):
                        return True
            visited[s_id] = 2  # visited
            return False

        for s_id in step_ids:
            if visited[s_id] == 0:
                dfs(s_id, [s_id])

        # 4. Helper to find all ancestors of a step
        def get_all_ancestors(s_id: str, memo: Dict[str, Set[str]]) -> Set[str]:
            if s_id in memo:
                return memo[s_id]
            ancestors = set()
            step = steps_dict.get(s_id)
            if step:
                for dep in step.depends_on:
                    ancestors.add(dep)
                    ancestors.update(get_all_ancestors(dep, memo))
            memo[s_id] = ancestors
            return ancestors

        ancestors_memo = {}
        
        # 5. Data flow validation (regex check for $step_id.field references)
        ref_pattern = re.compile(r"^\$(\w+)(?:\.(\w+))?$")
        
        def scan_for_refs(val: Any) -> List[tuple]:
            refs = []
            if isinstance(val, str):
                match = ref_pattern.match(val)
                if match:
                    refs.append((match.group(1), match.group(2)))
            elif isinstance(val, dict):
                for k, v in val.items():
                    refs.extend(scan_for_refs(v))
            elif isinstance(val, list):
                for item in val:
                    refs.extend(scan_for_refs(item))
            return refs

        for step in spec.steps:
            refs = scan_for_refs(step.input)
            for ref_step_id, ref_field in refs:
                if ref_step_id not in steps_dict:
                    errors.append(ValidationError(
                        code="INVALID_DATA_FLOW",
                        message=f"Step '{step.id}' references input from non-existent step '{ref_step_id}'",
                        path=f"steps.{step.id}.input"
                    ))
                    continue
                
                ancestors = get_all_ancestors(step.id, ancestors_memo)
                if ref_step_id != step.id and ref_step_id not in ancestors:
                    errors.append(ValidationError(
                        code="INVALID_DATA_FLOW",
                        message=f"Step '{step.id}' references output of step '{ref_step_id}' but does not depend on it",
                        path=f"steps.{step.id}.depends_on"
                    ))
                
                ref_step = steps_dict[ref_step_id]
                if ref_field and ref_step.output:
                    if ref_field not in ref_step.output:
                        errors.append(ValidationError(
                            code="INVALID_DATA_FLOW",
                            message=f"Step '{step.id}' references field '{ref_field}' which is not defined in step '{ref_step_id}' outputs",
                            path=f"steps.{step.id}.input"
                        ))

        return ValidationResult(valid=len(errors) == 0, errors=errors)