app.py

import asyncio
import aiohttp
import requests
import re
import pprint
import logging
from typing import List, Dict, Any, Union, Optional
from datetime import datetime
import gradio as gr
import asyncio
import pandas as pd

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# --- Configuration Constants ---
API_BASE_URL = "https://www.railyatri.in/get-next-days-sa-data"
TRAINS_API_URL = "https://trainticketapi.railyatri.in/api/trains-between-station-with-sa.json"
DEFAULT_QUOTA = "GN"
REQUEST_TIMEOUT = 30  # seconds
MAX_RETRIES = 3
RETRY_DELAY = 1  # seconds

# Valid quota codes that users can specify
VALID_QUOTAS = {
    "GN": "General",
    "TQ": "Tatkal",
    "PT": "Premium Tatkal", 
    "LD": "Ladies",
    "HP": "Handicapped",
    "YU": "Yuva",
    "SS": "Senior Citizen",
    "DF": "Defence",
    "RE": "Railway Employee",
    "RQ": "RAC"
}

# --- Custom Exceptions ---
class TrainAPIError(Exception):
    """Custom exception for train API related errors."""
    pass

class ValidationError(Exception):
    """Custom exception for input validation errors."""
    pass

# --- Validation Functions ---
def validate_station_code(station_code: str, field_name: str) -> str:
    """Validate and normalize station code."""
    if not isinstance(station_code, str):
        raise ValidationError(f"{field_name} must be a string")
    
    station = station_code.strip().upper()
    if not station or len(station) < 2:
        raise ValidationError(f"Invalid {field_name}: {station_code}")
    
    return station

def validate_date_format(date_str: str, input_format: str = None) -> str:
    """
    Validate date and convert to YYYY-M-D format.
    
    Args:
        date_str: Date string to validate
        input_format: Expected input format ('dd-mm-yyyy' or 'yyyy-mm-dd')
    
    Returns:
        str: Date in YYYY-M-D format
    """
    if not isinstance(date_str, str):
        raise ValidationError("Date must be a string")
    
    date_str = date_str.strip()
    
    # Try different date formats
    formats_to_try = []
    if input_format == 'dd-mm-yyyy':
        formats_to_try = ['%d-%m-%Y', '%d-%m-%y']
    elif input_format == 'yyyy-mm-dd':
        formats_to_try = ['%Y-%m-%d', '%Y-%M-%d']
    else:
        # Auto-detect format
        formats_to_try = ['%d-%m-%Y', '%Y-%m-%d', '%d-%m-%y', '%Y-%M-%d']
    
    for fmt in formats_to_try:
        try:
            dt_obj = datetime.strptime(date_str, fmt)
            return dt_obj.strftime('%Y-%m-%d')
        except ValueError:
            continue
    
    raise ValidationError(f"Invalid date format: {date_str}")

def validate_quota(quota: str) -> str:
    """Validate quota code."""
    if not isinstance(quota, str):
        raise ValidationError("Quota must be a string")
    
    quota = quota.strip().upper()
    if quota not in VALID_QUOTAS:
        valid_codes = ', '.join(VALID_QUOTAS.keys())
        raise ValidationError(f"Invalid quota '{quota}'. Valid codes: {valid_codes}")
    
    return quota

def validate_query(query: Dict[str, Any]) -> None:
    """Validates a single query dictionary for required fields and data types."""
    required_fields = ["train_number", "journey_class", "source", "destination", "journey_date"]
    
    for field in required_fields:
        if field not in query or query[field] is None:
            raise ValidationError(f"Missing required field: {field}")
    
    # Validate train_number
    train_num = str(query["train_number"]).strip()
    if not train_num or not train_num.isdigit():
        raise ValidationError(f"Invalid train_number: {query['train_number']}")
    
    # Validate journey_class
    journey_class = query["journey_class"]
    if isinstance(journey_class, str):
        journey_class = [journey_class]
    elif not isinstance(journey_class, list):
        raise ValidationError(f"journey_class must be string or list, got {type(journey_class)}")
    
    for j_class in journey_class:
        if not isinstance(j_class, str) or not j_class.strip():
            raise ValidationError(f"Invalid journey_class: {j_class}")
    
    # Validate station codes
    query["source"] = validate_station_code(query["source"], "source")
    query["destination"] = validate_station_code(query["destination"], "destination")
    
    # Validate date
    query["journey_date"] = validate_date_format(query["journey_date"], 'yyyy-mm-dd')
    
    # Validate quota if provided
    if "journey_quota" in query and query["journey_quota"] is not None:
        query["journey_quota"] = validate_quota(query["journey_quota"])

def validate_queries(queries: List[Dict[str, Any]]) -> None:
    """Validates a list of query dictionaries."""
    if not isinstance(queries, list):
        raise ValidationError("queries must be a list")
    
    if not queries:
        raise ValidationError("queries list cannot be empty")
    
    for i, query in enumerate(queries):
        if not isinstance(query, dict):
            raise ValidationError(f"Query at index {i} must be a dictionary")
        
        try:
            validate_query(query)
        except ValidationError as e:
            raise ValidationError(f"Query at index {i}: {str(e)}")

# --- Train Discovery Functions ---
def get_trains_between_stations(source: str, destination: str, journey_date: str) -> Optional[Dict[str, Any]]:
    """
    Fetch train details between two stations from Railyatri API.

    Args:
        source: Source station code (e.g., 'HWH').
        destination: Destination station code (e.g., 'PNBE').
        journey_date: Journey date in 'dd-mm-yyyy' format (e.g., '22-09-2025').

    Returns:
        Optional[Dict]: JSON response from API if successful, else None.
    """
    try:
        # Validate inputs
        source = validate_station_code(source, "source")
        destination = validate_station_code(destination, "destination")
        # Keep original date format for this API
        if not re.match(r'\d{2}-\d{2}-\d{4}', journey_date.strip()):
            # Try to convert from YYYY-MM-DD to DD-MM-YYYY if needed
            try:
                dt = datetime.strptime(journey_date.strip(), '%Y-%m-%d')
                journey_date = dt.strftime('%d-%m-%Y')
            except ValueError:
                raise ValidationError(f"Invalid date format: {journey_date}")
        
        params = {
            "from": source,
            "to": destination,
            "dateOfJourney": journey_date.strip()
        }
        
        logger.info(f"Fetching trains from {source} to {destination} on {journey_date}")
        
        response = requests.get(TRAINS_API_URL, params=params, timeout=REQUEST_TIMEOUT)
        response.raise_for_status()
        
        result = response.json()
        logger.info("Successfully fetched train list")
        return result
        
    except requests.exceptions.Timeout:
        logger.error("Timeout while fetching train list")
        return None
    except requests.exceptions.RequestException as e:
        logger.error(f"Error during train list request: {e}")
        return None
    except ValidationError as e:
        logger.error(f"Validation error: {e}")
        return None
    except Exception as e:
        logger.error(f"Unexpected error while fetching trains: {e}")
        return None

def simplify_train_info(api_result: Dict[str, Any], default_quota: str = DEFAULT_QUOTA) -> List[Dict[str, Any]]:
    """
    Simplify train information from API response.
    
    Args:
        api_result: Raw API response
        default_quota: Default quota to use for all trains
        
    Returns:
        List[Dict]: Simplified train information
    """
    if not isinstance(api_result, dict):
        logger.warning("Invalid API result format")
        return []
    
    simplified_trains = []

    def process_train_list(train_list: List[Dict[str, Any]]) -> None:
        """Process a list of trains and add to simplified_trains."""
        if not isinstance(train_list, list):
            return
            
        for train in train_list:
            if not isinstance(train, dict):
                continue
                
            try:
                # Handle date conversion
                raw_date = train.get("train_date", "").strip()
                formatted_date = None
                
                if raw_date:
                    # Try multiple date formats
                    for date_format in ['%d-%m-%Y', '%Y-%m-%d']:
                        try:
                            dt_obj = datetime.strptime(raw_date, date_format)
                            formatted_date = dt_obj.strftime('%Y-%m-%d')
                            break
                        except ValueError:
                            continue
                    
                    if not formatted_date:
                        logger.warning(f"Could not parse date: {raw_date}")
                        formatted_date = raw_date

                # Get journey classes, ensure it's a list
                journey_classes = train.get("journey_class", [])
                if isinstance(journey_classes, str):
                    journey_classes = [journey_classes]
                elif not isinstance(journey_classes, list):
                    journey_classes = []
                
                # Filter out empty classes
                journey_classes = [jc.strip() for jc in journey_classes if jc and str(jc).strip()]
                
                if not journey_classes:
                    logger.warning(f"No valid journey classes for train {train.get('train_number')}")
                    continue

                train_data = {
                    "train_number": str(train.get("train_number", "")).strip(),
                    "journey_class": journey_classes,
                    "source": str(train.get("from", "")).strip().upper(),
                    "destination": str(train.get("to", "")).strip().upper(),
                    "journey_date": formatted_date,
                    "journey_quota": validate_quota(default_quota)
                }
                
                # Validate essential fields
                if (train_data["train_number"] and 
                    train_data["source"] and 
                    train_data["destination"] and
                    train_data["journey_date"]):
                    simplified_trains.append(train_data)
                else:
                    logger.warning(f"Incomplete train data: {train_data}")
                    
            except Exception as e:
                logger.warning(f"Error processing train: {e}")
                continue

    # Process different train categories
    train_categories = [
        "train_between_stations",
        "alternate_trains", 
        "reserved_trains"
    ]
    
    for category in train_categories:
        if category in api_result:
            process_train_list(api_result[category])
    
    logger.info(f"Simplified {len(simplified_trains)} trains")
    return simplified_trains

# --- Core API and Parsing Functions ---
async def get_train_availability_async(
    session: aiohttp.ClientSession,
    query: Dict[str, Union[str, int]],
    retry_count: int = 0
) -> Dict[str, Any]:
    """
    Asynchronously fetches seat availability data from Railyatri for a single request.
    """
    # Normalize parameters
    params = {
        "train_number": str(query["train_number"]).strip(),
        "journey_class": str(query["journey_class"]).strip(),
        "from": str(query["source"]).strip().upper(),
        "to": str(query["destination"]).strip().upper(),
        "journey_date": str(query["journey_date"]).strip(),
        "journey_quota": str(query.get("journey_quota", DEFAULT_QUOTA)).strip(),
    }
    
    try:
        timeout = aiohttp.ClientTimeout(total=REQUEST_TIMEOUT)
        async with session.get(API_BASE_URL, params=params, timeout=timeout) as response:
            response.raise_for_status()
            
            content_type = response.headers.get('content-type', '').lower()
            if 'application/json' not in content_type:
                logger.warning(f"Unexpected content type for train {params['train_number']}")
            
            response_data = await response.json()
            
            if not isinstance(response_data, dict):
                logger.warning(f"Invalid response format for train {params['train_number']}")
                return {}
            
            return response_data
            
    except asyncio.TimeoutError:
        logger.error(f"Timeout for train {params['train_number']} class {params['journey_class']}")
        if retry_count < MAX_RETRIES:
            await asyncio.sleep(RETRY_DELAY * (retry_count + 1))
            return await get_train_availability_async(session, query, retry_count + 1)
        return {}
        
    except aiohttp.ClientResponseError as e:
        if e.status == 429 and retry_count < MAX_RETRIES:
            await asyncio.sleep(RETRY_DELAY * (retry_count + 1) * 2)
            return await get_train_availability_async(session, query, retry_count + 1)
        logger.error(f"HTTP {e.status} for train {params['train_number']}: {e}")
        return {}
        
    except Exception as e:
        logger.error(f"Error for train {params['train_number']}: {e}")
        if retry_count < MAX_RETRIES:
            await asyncio.sleep(RETRY_DELAY * (retry_count + 1))
            return await get_train_availability_async(session, query, retry_count + 1)
        return {}

def safe_extract_train_number(ticket_link: str) -> str:
    """Safely extracts train number from ticket link."""
    if not isinstance(ticket_link, str):
        return "N/A"
    
    try:
        match = re.search(r'train_no=(\d+)', ticket_link)
        return match.group(1) if match else "N/A"
    except Exception:
        return "N/A"

def extract_train_availability(data: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Extracts and transforms train availability details from raw API data."""
    if not isinstance(data, dict):
        return []
    
    extracted_data = []
    
    try:
        # Find seat availability data with multiple fallback patterns
        seat_availability = None
        
        # Check nested structure first
        if "result" in data and isinstance(data["result"], dict):
            for key in ["seat_availibility", "seat_availability"]:
                if key in data["result"] and isinstance(data["result"][key], list):
                    seat_availability = data["result"][key]
                    break
        
        # Check root level
        if not seat_availability:
            for key in ["seat_availibility", "seat_availability"]:
                if key in data and isinstance(data[key], list):
                    seat_availability = data[key]
                    break
        
        if not seat_availability:
            return []

        for entry in seat_availability:
            if not isinstance(entry, dict):
                continue
            
            # Extract fields with fallbacks
            ticket_link = entry.get("ticket_link", "")
            train_no = safe_extract_train_number(ticket_link)
            
            class_type = None
            for key in ["class_type", "class", "journey_class"]:
                if key in entry and entry[key]:
                    class_type = str(entry[key]).strip()
                    break
            
            if not class_type:
                continue
            
            # Get availability
            availability = entry.get(class_type)
            if availability is None:
                for key in ["availability", "status", "seats"]:
                    if key in entry:
                        availability = entry[key]
                        break
            
            # Get date
            journey_date = entry.get("Date (DD-MM-YYYY)")
            if journey_date is None:
                for key in ["date", "journey_date", "Date"]:
                    if key in entry:
                        journey_date = entry[key]
                        break
            
            # Get fare
            total_fare = entry.get("total_fare")
            if total_fare is None:
                for key in ["fare", "price", "cost"]:
                    if key in entry:
                        total_fare = entry[key]
                        break
            
            extracted_data.append({
                "train_no": train_no,
                "journey_date": journey_date,
                "class_type": class_type,
                "availability": availability,
                "total_fare": total_fare
            })
            
    except Exception as e:
        logger.error(f"Error extracting availability data: {e}")
        return []

    return extracted_data

# --- Main Functions ---
async def fetch_all_availabilities(queries: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    Asynchronously fetches and processes train availability for multiple queries.
    
    Args:
        queries: List of query dictionaries
        
    Returns:
        List of extracted availability information
    """
    try:
        validate_queries(queries)
    except ValidationError as e:
        logger.error(f"Input validation failed: {e}")
        raise
    
    logger.info(f"Processing {len(queries)} queries")
    
    tasks = []
    connector = aiohttp.TCPConnector(limit=10, limit_per_host=5)
    timeout = aiohttp.ClientTimeout(total=REQUEST_TIMEOUT)
    
    try:
        async with aiohttp.ClientSession(
            connector=connector,
            timeout=timeout,
            headers={'User-Agent': 'Python-TrainAvailability/1.0'}
        ) as session:
            
            for query in queries:
                journey_classes = query["journey_class"]
                if isinstance(journey_classes, str):
                    journey_classes = [journey_classes.strip()]
                elif isinstance(journey_classes, list):
                    journey_classes = [str(jc).strip() for jc in journey_classes if jc]

                for j_class in journey_classes:
                    if not j_class:
                        continue
                        
                    single_request_query = query.copy()
                    single_request_query["journey_class"] = j_class
                    
                    task = get_train_availability_async(session, single_request_query)
                    tasks.append(task)
            
            if not tasks:
                return []
            
            api_responses = await asyncio.gather(*tasks, return_exceptions=True)
        
    except Exception as e:
        logger.error(f"Error executing requests: {e}")
        raise TrainAPIError(f"Failed to execute API requests: {e}")
    
    # Process results
    all_availability_data = []
    successful_requests = 0
    
    for response in api_responses:
        if isinstance(response, Exception) or not isinstance(response, dict) or not response:
            continue
            
        try:
            extracted_info = extract_train_availability(response)
            if extracted_info:
                all_availability_data.extend(extracted_info)
                successful_requests += 1
        except Exception as e:
            logger.error(f"Error processing response: {e}")
            continue
    
    logger.info(f"Processed {successful_requests}/{len(tasks)} requests successfully")
    logger.info(f"Extracted {len(all_availability_data)} availability records")
    
    return all_availability_data
@gr.mcp.tool()
async def find_available_seats(
    source: str, 
    destination: str, 
    journey_date: str, 
    quota: str = DEFAULT_QUOTA
) -> List[Dict[str, Any]]:
    """
    Complete workflow to find available seats between stations.
    
    Args:
        source: Source station code (e.g., 'HWH')
        destination: Destination station code (e.g., 'PNBE') 
        journey_date: Journey date in 'dd-mm-yyyy' format (e.g., '22-09-2025')
        quota: Journey quota (default: 'GN')
        
    Returns:
        List of available seat information
        
    Raises:
        ValidationError: If input validation fails
        TrainAPIError: If API requests fail
    """
    try:
        # Validate inputs
        source = validate_station_code(source, "source")
        destination = validate_station_code(destination, "destination") 
        quota = validate_quota(quota)
        
        logger.info(f"Finding available seats: {source} → {destination} on {journey_date} (quota: {quota})")
        
        # Step 1: Get trains between stations
        trains_data = get_trains_between_stations(source, destination, journey_date)
        if not trains_data:
            logger.warning("No trains data received")
            return []
        
        # Step 2: Simplify train information
        simplified_trains = simplify_train_info(trains_data, quota)
        if not simplified_trains:
            logger.warning("No valid trains found")
            return []
        
        logger.info(f"Found {len(simplified_trains)} trains to check")
        
        # Step 3: Fetch availability for all trains
        availability_data = await fetch_all_availabilities(simplified_trains)
        
        logger.info(f"Seat search completed: found {len(availability_data)} availability records")
        return availability_data
        
    except ValidationError:
        raise
    except TrainAPIError:
        raise
    except Exception as e:
        logger.error(f"Unexpected error in find_available_seats: {e}")
        raise TrainAPIError(f"Unexpected error: {e}")

# --- Gradio Wrapper ---
def gradio_find_seats(source: str, destination: str, journey_date: str, quota: str):
    try:
        # Run async function in sync wrapper
        result = asyncio.run(find_available_seats(source, destination, journey_date, quota))
        
        if not result:
            return "No availability found", None
        
        # Convert to DataFrame for nice table display
        df = pd.DataFrame(result)
        return "Success", df
    except Exception as e:
        return f"Error: {e}", None

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("## 🚆 Train Seat Availability Finder")
    
    with gr.Row():
        source = gr.Textbox(label="Source Station Code", value="HWH")
        destination = gr.Textbox(label="Destination Station Code", value="PNBE")
    
    with gr.Row():
        journey_date = gr.Textbox(label="Journey Date (dd-mm-yyyy)", value="22-09-2025")
        quota = gr.Dropdown(choices=list(VALID_QUOTAS.keys()), value="GN", label="Quota")
    
    submit_btn = gr.Button("Find Seats")
    
    status = gr.Textbox(label="Status", interactive=False)
    output_table = gr.Dataframe(label="Availability", interactive=False)
    
    submit_btn.click(
        fn=gradio_find_seats,
        inputs=[source, destination, journey_date, quota],
        outputs=[status, output_table]
    )

# Launch app
if __name__ == "__main__":
    demo.launch(mcp_server= True)