from typing import List, Optional, Dict, Union
from sqlalchemy.orm import Session
from sqlalchemy import func
from models.country import Country
from pydantic import BaseModel, constr

def validate_country_code(country_code: str) -> bool:
    """
    Validate that country code is exactly 2 characters and contains only letters.
    
    Args:
        country_code: The country code to validate
        
    Returns:
        bool: True if valid, False otherwise
    """
    if len(country_code) != 2:
        return False
    return country_code.isalpha()

def search_countries(db: Session, search_query: Optional[str] = None) -> List[Country]:
    """
    Search for countries that start with the given search query.
    
    Args:
        db: Database session
        search_query: Optional search string to filter country names
        
    Returns:
        List of matching Country objects
    """
    query = db.query(Country)
    
    if search_query:
        query = query.filter(Country.name.ilike(f'{search_query}%'))
    
    return query.order_by(Country.name).all()

def get_country_by_code(db: Session, country_code: str) -> Optional[Country]:
    """
    Get a country by its unique code.
    
    Args:
        db: Database session
        country_code: Two letter country code
        
    Returns:
        Country object if found, None otherwise
    """
    return db.query(Country).filter(Country.code == country_code.upper()).first()

def create_country_safely(db: Session, name: str, code: str) -> Union[Country, Dict[str, str]]:
    """
    Create a new country with validation and error handling.
    
    Args:
        db: Database session
        name: Country name
        code: Country code
        
    Returns:
        Country object if created successfully, error dict otherwise
    """
    if not validate_country_code(code):
        return {"error": "Invalid country code format"}
        
    existing_country = db.query(Country).filter(
        func.lower(Country.name) == name.lower()
    ).first()
    
    if existing_country:
        return {"error": "Country name already exists"}
        
    existing_code = get_country_by_code(db, code)
    if existing_code:
        return {"error": "Country code already exists"}
        
    db_country = Country(
        name=name,
        code=code.upper()
    )
    
    db.add(db_country)
    db.commit()
    db.refresh(db_country)
    
    return db_country

def format_country_response(country: Country) -> Dict[str, str]:
    """
    Format a country object into a standardized response format.
    
    Args:
        country: Country object to format
        
    Returns:
        Dict containing formatted country data
    """
    return {
        "name": country.name,
        "code": country.code,
    }