import inspect 
import random 
from typing import TypedDict 

from crbench .base import Instance ,Program ,Task 

GraphReachabilityInput =TypedDict (
"GraphReachabilityInput",
{"num_vertices":int ,"edges":list [tuple [int ,int ]],"source":int ,"target":int },
)


def is_reachable (
num_vertices :int ,
edges :list [tuple [int ,int ]],
source :int ,
target :int ,
)->bool :


    graph ={u :[]for u in range (num_vertices )}
    for u ,v in edges :
        graph [u ].append (v )

    visited =set ()
    stack =[source ]
    while stack :
        u =stack .pop ()
        if u ==target :
            return True 
        if u in visited :
            continue 
        visited .add (u )
        stack .extend (v for v in graph [u ]if v not in visited )

    return False 


class GraphReachabilityTask (Task [GraphReachabilityInput ,bool ,bool ]):
    def __init__ (
    self ,n_range :tuple [int ,int ]=(4 ,10 ),m_range :tuple [int ,int ]=(10 ,20 )
    ):
        super ().__init__ ()
        self .n_range =n_range 
        self .m_range =m_range 

    def get_instance (self ,seed :int )->Instance [GraphReachabilityInput ,bool ]:
        rng =random .Random (seed )


        expected_answer =rng .choice ([True ,False ])

        while True :
            n =rng .randint (self .n_range [0 ],self .n_range [1 ])
            m =rng .randint (self .m_range [0 ],min (self .m_range [1 ],n *(n -1 )))

            possible_edges =[(i ,j )for i in range (n )for j in range (n )if i !=j ]
            edge_list =rng .sample (possible_edges ,m )
            rng .shuffle (edge_list )

            while True :
                source =rng .randint (0 ,n -1 )
                target =rng .randint (0 ,n -1 )
                if source !=target :
                    break 

            input :GraphReachabilityInput ={
            "num_vertices":n ,
            "edges":edge_list ,
            "source":source ,
            "target":target ,
            }
            answer =is_reachable (**input )
            if answer ==expected_answer :
                break 

        return Instance (
        input =input ,
        verifier_hint =is_reachable (**input ),
        )

    def get_solution_program (self )->Program :
        return Program (
        source_code =inspect .getsource (is_reachable ),
        entry_point ="is_reachable",
        )

    def solve (self ,input :GraphReachabilityInput )->bool :
        return is_reachable (**input )

    def verify (
    self ,instance :Instance [GraphReachabilityInput ,bool ],output :bool 
    )->bool :
        return output ==instance .verifier_hint 

    def verify_from_str (
    self ,instance :Instance [GraphReachabilityInput ,bool ],output_str :str 
    )->bool :
        if output_str not in ["True","False"]:
            return False 
        output =bool (output_str =="True")
        return self .verify (instance ,output )

    def build_prompt (self ,instance :Instance [GraphReachabilityInput ,bool ])->str :
        edges_str ="\n".join (f"{u } -> {v }"for u ,v in instance .input ["edges"])

        return f"""You are given a directional graph with {instance .input ["num_vertices"]} vertices labeled from 0 to {instance .input ["num_vertices"]-1 }. Is it possible to reach from vertex {instance .input ["source"]} to vertex {instance .input ["target"]}? The edges are given as follows:

```
{edges_str }
```
The answer should be True or False.
"""


class GraphReachabilityTaskSmall (GraphReachabilityTask ):
    def __init__ (self )->None :
        super ().__init__ (n_range =(4 ,8 ),m_range =(1 ,15 ))


class GraphReachabilityTaskMedium (GraphReachabilityTask ):
    def __init__ (self )->None :
        super ().__init__ (n_range =(50 ,100 ),m_range =(50 ,150 ))


class GraphReachabilityTaskLarge (GraphReachabilityTask ):
    def __init__ (self )->None :
        super ().__init__ (n_range =(150 ,200 ),m_range =(200 ,500 ))
