from datetime import datetime
from pydantic import BaseModel, TypeAdapter
from typing_extensions import Iterable, Optional, Any, Union, List
import json
import pytz
import os
import warnings
import requests
from urllib.parse import urljoin
from pathlib import Path
API_TOKEN_ENV = "ETL_API_TOKEN"
from etlup import TestModel, TestType, SequenceModel, SequenceType, Tests
from .sequences import BaseSequence
TestArrModel = TypeAdapter(List[TestType])

class Config:
    """Configuration class for ETL session settings."""
    
    def __init__(self):
        self.api_token: Optional[str] = None

    def load_from_env(self, dotenv_path: Union[None, str]):
        """Load configuration from environment variables and .env file."""
        try:
            from dotenv import load_dotenv, find_dotenv
            if dotenv_path is None:
                dotenv_path = find_dotenv(usecwd=True)

            if not Path(dotenv_path).is_file():
                raise FileNotFoundError(f"Your .env file at {dotenv_path} does not exist.")
            
            env_loaded = load_dotenv(dotenv_path, override=True)
            print(f"loaded env: {env_loaded}")
        except ImportError:
            pass
        
        self.api_token = os.getenv(API_TOKEN_ENV)
        return self
    
    def validate_api_token(self) -> str:
        """Validate and return the API token, raising error if not found."""
        if self.api_token and self.api_token.strip():
            return self.api_token.strip()
        raise ValueError(
            f"API token not found. Set {API_TOKEN_ENV}=<token> in your environment "
            f"or .env file."
        )

def user_default_tz():
    """
    !NOT IMPLEMENTED! It would be nice to auto catch user timezone from their computer maybe
    """
    user_tz = datetime.now().astimezone().tzinfo
    warnings.warn(f"Explicit timezone information was not provided using your system default: {user_tz}")
    return user_tz

def localize_datetime(dt: datetime, tz: str = "auto") -> datetime:
    """
    Localize or convert a datetime to the specified timezone.

    - tz='auto': use the computer's local timezone.
    - If dt is naive: assign the timezone (pytz.localize when available).
    - If dt is aware: convert to the target timezone.
    """
    if not isinstance(dt, datetime):
        raise TypeError("dt must be a datetime.datetime")

    if tz == "auto":
        target_tz = datetime.now().astimezone().tzinfo  # local system tz
    else:
        if tz not in pytz.all_timezones:
            raise ValueError(
                f"Timezone '{tz}' is not valid. See: https://en.wikipedia.org/wiki/List_of_tz_database_time_zones"
            )
        target_tz = pytz.timezone(tz)

    if dt.tzinfo is None:
        # naive -> assign timezone
        localizer = getattr(target_tz, "localize", None)
        return localizer(dt) if callable(localizer) else dt.replace(tzinfo=target_tz)

    # aware -> convert
    return dt.astimezone(target_tz)

def now_utc() -> datetime:
    """
    Return the current time as a timezone-aware UTC datetime.
    """
    return datetime.now(tz=pytz.UTC)

def get_model(version: str, constr_type: str):
    def _literal_values(prop: dict) -> list:
        """Return the list of allowed values for a Literal schema property."""
        if 'const' in prop:
            return [prop['const']]
        return prop["enum"]

    for test_cls in Tests:
        props = test_cls.model_json_schema()['properties']

        names = _literal_values(props['name'])
        versions = _literal_values(props['version'])

        if constr_type in names and version in versions:
            return test_cls

class Session:
    """
    A lightweight API for committing ETL assemblies and tests.

    Options:
    prod: bool, If true it will be uploaded to the production instance
    domain: str, Overwrites prod and will attempt to upload to this domain, exmaples: http://127.0.0.1/ or https://prod-etl.app.cern.ch/
    for_assembly: bool, If true it uploads it as an assembly, otherwise a test, default is False
    
    Usage:
    from etlup import now_utc, localize_datetime, Session, tamalero
    import numpy as np

    bl = tamalero.BaselineV0(
        module = "PBU003",
        location = "BU",
        user_created = "hswanson",
        measurement_date = now_utc(),
        "U1": np.zeros((16,16)),
        "U2": np.ones((16,16)),
        "U3": np.zeros((16,16)),
        "U4": np.ones((16,16))
    )
    with Session(domain="http://127.0.0.1/") as sesh:    
        sesh.add(bl)
        sesh.to_file()
        sesh.upload()
    """

    def __init__(self, prod=False, domain=None, for_assembly=False):
        self._constrs: List[BaseModel] = []  # only Pydantic models
        self._sequence: Optional[SequenceType] = None
        if domain is not None:
            self.api_domain = domain
        elif prod == True:
            self.api_domain = "https://prod-etl.app.cern.ch/"
        else:
            self.api_domain = "https://staging-etl.app.cern.ch/"
        
        endpoint = f"api/upload/{'assembly' if for_assembly else 'test'}"
        base = self.api_domain.rstrip('/') + '/'
        self.api_url = urljoin(base, endpoint.lstrip('/'))
        self.seq_url = urljoin(base, "api/upload/sequence")
        #print(f"Using the API URL: {self.api_url}")

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        # Always clear pending constructions on exit
        self.clear()
        # Propagate any exception
        return False

    def add(self, constr: BaseModel):
        """
        Add a single construction payload (Pydantic model only).
        """
        if not isinstance(constr, BaseModel):
            raise TypeError("Only Pydantic BaseModel instances can be added. Dicts are not supported.")
        
        if isinstance(constr, BaseSequence):
            if self._sequence is not None:
                raise ValueError("Only one sequence can be loaded into the session at a time. Please upload or clear before adding another one.")
            self._sequence = constr
        else:
            self._constrs.append(constr)
        return self

    def add_all(self, items: Iterable[BaseModel]):
        for item in items:
            self.add(item)
        return self

    def clear_sequence(self):
        self._sequence = None
    
    def clear(self):
        """
        Removes all the loaded assemblies or tests
        """
        self._constrs.clear()
        self.clear_sequence()

    def _build_payload(self) -> Any:
        """
        Build a JSON-serializable payload for the pending constructions.
        Uses the TestArrModel TypeAdapter to ensure schema correctness.
        """
        if not self._constrs:
            raise ValueError("Cannot commit when no assembly or tests were added")
        
        # dump_python(..., mode='json') returns JSON-ready Python data
        return TestArrModel.dump_python(self._constrs, mode="json")

    def to_file(self, filepath: str = "") -> Union[str, List[str]]:
        """
        Serialize the pending constructions to a JSON file (no network).
        """
        files_created = []

        if self._constrs:
            if not filepath:
                test_filepath = Path.cwd() / 'construction_upload.json'
            else:
                test_filepath = Path(filepath)

            payload = self._build_payload()
            with open(test_filepath, 'w') as json_file:
                json.dump(payload, json_file, indent=4)
            files_created.append(str(test_filepath))
        
        if self._sequence:
            out_dir = Path(filepath).parent if filepath else Path.cwd()
            filename = "construction_sequence.json"
            seq_path = out_dir / filename
            
            with open(seq_path, 'w') as f:
                json.dump(self._sequence.model_dump(mode='json'), f, indent=4)
            files_created.append(str(seq_path))

        if not files_created:
            raise ValueError("Cannot write to file when no assembly, tests, or sequences were added")
            
        if len(files_created) == 1:
            return files_created[0]
        return files_created

    def upload(
        self,
        timeout: float = 30.0,
        dry_run: bool = False,
        dotenv_path: str = None
    ) -> dict:
        """
        Commit pending constructions to the API.

        - base_url: overrides ETL_API_URL
        - dry_run: if True, do not POST; return the payload instead
        """
        self.config = Config().load_from_env(dotenv_path)

        # Prepare payloads
        test_payload = None
        
        if self._constrs:
            test_payload = self._build_payload()

        if not test_payload and not self._sequence:
            raise ValueError("Nothing to upload!")

        if dry_run:
            seq_payload = self._sequence.model_dump(mode='json') if self._sequence else None
            count = len(self._constrs)
            if self._sequence:
                count += 1
            return {
                "dry_run": True, 
                "test_payload": test_payload, 
                "seq_payload": seq_payload,
                "count": count
            }

        auth_token = self.config.validate_api_token()
        headers = {
            "Authorization": f"Bearer {auth_token}",
            "Content-Type": "application/json",
        }
        
        responses = {}

        if test_payload:
            resp = requests.post(self.api_url, 
                                    json=test_payload, 
                                    headers=headers, 
                                    timeout=timeout)
            try:
                resp.raise_for_status()
            except requests.HTTPError as exc:
                data = resp.text
                raise requests.HTTPError(
                    f"{data}",
                    response = resp,
                    request = exc.request
                )
            try:
                responses['test_response'] = resp.json()
            except ValueError:
                raise RuntimeError("Upload succeeded but response body was not valid JSON")

        if self._sequence:
            payload = self._sequence.model_dump(mode='json')
            resp = requests.post(self.seq_url, 
                                    json=payload, 
                                    headers=headers, 
                                    timeout=timeout)
            try:
                resp.raise_for_status()
            except requests.HTTPError as exc:
                data = resp.text
                raise requests.HTTPError(
                    f"{data}",
                    response = resp,
                    request = exc.request
                )
            try:
                responses['seq_response'] = resp.json()
            except ValueError:
                raise RuntimeError("Upload succeeded but response body was not valid JSON")
        
        self.clear()
        
        # If only one type was uploaded, return that response directly for backward compatibility
        if len(responses) == 1:
            return list(responses.values())[0]
            
        return responses